From ee2fd7ea186cfe5a24aede967267c7972e000981 Mon Sep 17 00:00:00 2001
From: Jianhui Zhao <jianhuizhao329@gmail.com>
Date: Thu, 25 Apr 2019 00:20:06 +0800
Subject: [PATCH] Support device authentication

Signed-off-by: Jianhui Zhao <jianhuizhao329@gmail.com>
---
 client.go  |  9 ++++++++-
 http.go    |  2 +-
 main.go    | 12 +++++++++++-
 rttys.conf | 18 ++++++++++--------
 4 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/client.go b/client.go
index d161dcf..25284bf 100644
--- a/client.go
+++ b/client.go
@@ -82,7 +82,14 @@ func (c *Client) writePump() {
 }
 
 /* serveWs handles websocket requests from the device or user. */
-func serveWs(br *Broker, w http.ResponseWriter, r *http.Request) {
+func serveWs(br *Broker, w http.ResponseWriter, r *http.Request, cfg *RttysConfig) {
+	token := r.Header.Get("Authorization")
+	if token != cfg.token {
+		log.Error("Invalid token from terminal device")
+		http.Error(w, "Forbidden", http.StatusForbidden)
+		return
+	}
+
 	keepalive, _ := strconv.Atoi(r.URL.Query().Get("keepalive"))
 	isDev := r.URL.Query().Get("device") != ""
 	devid := r.URL.Query().Get("devid")
diff --git a/http.go b/http.go
index cc8da27..d5ebb37 100644
--- a/http.go
+++ b/http.go
@@ -86,7 +86,7 @@ func httpStart(br *Broker, cfg *RttysConfig) {
 	staticfs := http.FileServer(statikFS)
 
 	http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
-		serveWs(br, w, r)
+		serveWs(br, w, r, cfg)
 	})
 
 	http.HandleFunc("/cmd", func(w http.ResponseWriter, r *http.Request) {
diff --git a/main.go b/main.go
index 5a0d1d8..a299a9d 100644
--- a/main.go
+++ b/main.go
@@ -25,6 +25,7 @@ import (
 	"encoding/binary"
 	"encoding/hex"
 	"flag"
+	"fmt"
 	"io"
 	"os"
 	"runtime"
@@ -42,6 +43,7 @@ type RttysConfig struct {
 	sslKey   string
 	username string
 	password string
+	token    string
 }
 
 func init() {
@@ -79,7 +81,7 @@ func genUniqueID(extra string) string {
 }
 
 func setConfigOpt(yamlCfg *yaml.File, name string, opt *string) {
-	val, err := yamlCfg.Get("addr")
+	val, err := yamlCfg.Get(name)
 	if err != nil {
 		return
 	}
@@ -92,10 +94,17 @@ func parseConfig() *RttysConfig {
 	flag.StringVar(&cfg.addr, "addr", ":5912", "address to listen")
 	flag.StringVar(&cfg.sslCert, "ssl-cert", "./rttys.crt", "certFile Path")
 	flag.StringVar(&cfg.sslKey, "ssl-key", "./rttys.key", "keyFile Path")
+	flag.StringVar(&cfg.token, "token", "", "token to use")
 	conf := flag.String("conf", "./rttys.conf", "config file to load")
+	genToken := flag.Bool("gen-token", false, "generate token")
 
 	flag.Parse()
 
+	if *genToken {
+		fmt.Println(genUniqueID("rttys-token"))
+		os.Exit(0)
+	}
+
 	yamlCfg, err := yaml.ReadFile(*conf)
 	if err == nil {
 		setConfigOpt(yamlCfg, "addr", &cfg.addr)
@@ -103,6 +112,7 @@ func parseConfig() *RttysConfig {
 		setConfigOpt(yamlCfg, "ssl-key", &cfg.sslKey)
 		setConfigOpt(yamlCfg, "username", &cfg.username)
 		setConfigOpt(yamlCfg, "password", &cfg.password)
+		setConfigOpt(yamlCfg, "token", &cfg.token)
 	}
 
 	return cfg
diff --git a/rttys.conf b/rttys.conf
index a46d26b..f2ef9ec 100644
--- a/rttys.conf
+++ b/rttys.conf
@@ -1,8 +1,10 @@
-#addr: :5912
-
-# default from system
-#username: rttys
-#password: rttys
-
-#ssl-cert: /etc/rttys/rttys.crt
-#ssl-key: /etc/rttys/rttys.key
+#addr: :5912
+
+# default from system
+#username: rttys
+#password: rttys
+
+#ssl-cert: /etc/rttys/rttys.crt
+#ssl-key: /etc/rttys/rttys.key
+
+#token: a1d4cdb1a3cd6a0e94aa3599afcddcf5