unify auth management

auth/auth.go

@@ -52,6 +52,39 @@ func WsIsAuthenticated(conn *websocket.Conn, token string) bool {
 	}
 }
 
+// this takes a bool for auth/unauth -- also persists to config
+func AddToAuthMap(conn *websocket.Conn, token map[string]string, authed bool) error {
+	tokenStr := token["token"]
+	tokenId := token["id"]
+	hashed := sha256.Sum256([]byte(tokenStr))
+	hash := hex.EncodeToString(hashed[:])
+	if authed {
+		AuthenticatedClients.Lock()
+		AuthenticatedClients.Conns[tokenId] = conn
+		AuthenticatedClients.Unlock()
+		UnauthClients.Lock()
+		if _, ok := UnauthClients.Conns[tokenId]; ok {
+			delete(UnauthClients.Conns, tokenId)
+		}
+		UnauthClients.Unlock()
+	} else {
+		UnauthClients.Lock()
+		UnauthClients.Conns[tokenId] = conn
+		UnauthClients.Unlock()
+		AuthenticatedClients.Lock()
+		if _, ok := AuthenticatedClients.Conns[tokenId]; ok {
+			delete(AuthenticatedClients.Conns, tokenId)
+		}
+		AuthenticatedClients.Unlock()
+	}
+	now := time.Now().Format("2006-01-02_15:04:05")
+	err := AddSession(tokenId, hash, now, authed)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 // check the validity of the token
 func CheckToken(token map[string]string, conn *websocket.Conn, r *http.Request, setup bool) (bool, string, error) {
 	// great you have token. we see if valid.
@@ -119,14 +152,13 @@ func CreateToken(conn *websocket.Conn, r *http.Request, setup bool) (map[string]
 	if err != nil {
 		return nil, fmt.Errorf("failed to encrypt token: %v", err)
 	}
-	hashed := sha256.Sum256([]byte(encryptedText))
-	hash := hex.EncodeToString(hashed[:])
-	// Update sessions in the system's configuration
-	AddSession(id, hash, now, setup)
-	return map[string]string{
+	token := map[string]string{
 		"id":    id,
 		"token": encryptedText,
-	}, nil
+	}
+	// Update sessions in the system's configuration
+	AddToAuthMap(conn, token, setup)
+	return token, nil
 }
 
 // take session details and add to SysConfig

ws/ws.go

@@ -107,7 +107,7 @@ func WsHandler(w http.ResponseWriter, r *http.Request) {
 		payload.Payload = structs.WsLoginPayload{}
 		switch payload.Type {
 		case "login":
-			if err = loginHandler(msg, payload); err != nil {
+			if err = loginHandler(conn, msg, payload); err != nil {
 				config.Logger.Error(fmt.Sprintf("%v", err))
 			}
 		case "setup":
@@ -142,9 +142,8 @@ func WsHandler(w http.ResponseWriter, r *http.Request) {
 }
 
 // validate password and add to auth session map
-func loginHandler(msg []byte, payload structs.WsPayload) error {
+func loginHandler(conn *websocket.Conn, msg []byte, payload structs.WsPayload) error {
 	config.Logger.Info("Login")
-	now := time.Now().Format("2006-01-02_15:04:05")
 	payload.Payload = structs.WsLoginPayload{}
 	if err := json.Unmarshal(msg, &payload); err != nil {
 		return fmt.Errorf("Error unmarshalling message: %v", err)
@@ -155,7 +154,11 @@ func loginHandler(msg []byte, payload structs.WsPayload) error {
 	}
 	isAuthenticated := auth.AuthenticateLogin(loginPayload.Password)
 	if isAuthenticated {
-		if err := auth.AddSession(payload.Token.ID, payload.Token.Token, now, true); err != nil {
+		token := map[string]string{
+			"id":    payload.Token.ID,
+			"token": payload.Token.Token,
+		}
+		if err := auth.AddToAuthMap(conn, token, true); err != nil {
 			return fmt.Errorf("Unable to process login: %v", err)
 		}
 	} else {