goseg/ws/ws.go

package ws

import (
	"encoding/json"
	"fmt"
	"goseg/auth"
	"goseg/broadcast"
	"goseg/config"
	"goseg/structs"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
)

var (
	upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			return true // Allow all origins
		},
	}
)

// func handleConnection(c *websocket.Conn) {
//     // Read the first message from the client which should be the token
//     messageType, p, err := c.ReadMessage()
//     if err != nil {
//         config.Logger.Error(fmt.Errorf("%v",err))
//         return
//     }
//     token := string(p)
//     // Verify the token
//     isValid, _, err := CheckToken(token, c, false)  // 'false' assumes it's not a setup
//     if !isValid || err != nil {
//         config.Logger.Info("Invalid token provided by client.")
//         c.Close()
//         return
//     }

//     // rest of logic
// }

// switch on ws event cases
func WsHandler(w http.ResponseWriter, r *http.Request) {
	conf := config.Conf()
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		config.Logger.Error(fmt.Sprintf("Couldn't upgrade websocket connection: %v", err))
		return
	}
	// manage broadcasts and clients thru the broadcast package
	// broadcast.RegisterClient(conn)
	// defer broadcast.UnregisterClient(conn)
	if !auth.WsAuthCheck(conn) {
		unauthHandler(conn, r)
	}
	// keepalive for ws
	conn.SetPongHandler(func(string) error {
		conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})
	pingInterval := 15 * time.Second
	go func() {
		ticker := time.NewTicker(pingInterval)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
					return
				}
			}
		}
	}()
	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			return
		}
		var payload structs.WsPayload
		// if it fails to marshal, assume that means no token
		// "verify" action is implicit
		if err := json.Unmarshal(msg, &payload); err != nil {
			config.Logger.Info("New client")
			// auth.CreateToken also adds to unauth map
			newToken, err := auth.CreateToken(conn, r, false)
			if err != nil {
				config.Logger.Error("Unable to create token")
				continue
			}
			result := map[string]interface{}{
				"type":     "activity",
				"id":       payload.ID, // this is like the action id
				"error":    "null",
				"response": "ack",
				"token":    newToken,
			}
			respJson, err := json.Marshal(result)
			if err != nil {
				errmsg := fmt.Sprintf("Error marshalling token: %v", err)
				config.Logger.Error(errmsg)
				return
			}
			if err := conn.WriteMessage(websocket.TextMessage, respJson); err != nil {
				config.Logger.Error(fmt.Sprintf("Error writing response: %v", err))
				return
			}
		}
		payload.Payload = structs.WsLoginPayload{}
		token := map[string]string{
			"id":    payload.Token.ID,
			"token": payload.Token.Token,
		}
		if auth.CheckToken(token, conn, r, conf.FirstBoot) {
			switch payload.Type {
			case "new_ship":
				config.Logger.Info("New ship")
			case "pier_upload":
				config.Logger.Info("Pier upload")
			case "password":
				config.Logger.Info("Password")
			case "system":
				config.Logger.Info("System")
			case "startram":
				config.Logger.Info("StarTram")
			case "urbit":
				config.Logger.Info("Urbit")
			case "support":
				if err = supportHandler(msg, payload, r, conn); err != nil {
					config.Logger.Error(fmt.Sprintf("%v", err))
				}
			case "broadcast":
				if err := broadcast.BroadcastToClients(); err != nil {
					errmsg := fmt.Sprintf("Unable to broadcast to peer(s): %v", err)
					config.Logger.Error(errmsg)
					return
				}
			default:
				errmsg := fmt.Sprintf("Unknown request type: %s", payload.Type)
				config.Logger.Warn(errmsg)
				return
			}
		} else {
			switch payload.Type {
			case "login":
				if err = loginHandler(conn, msg, payload); err != nil {
					config.Logger.Error(fmt.Sprintf("%v", err))
					return
				}
			case "setup":
				config.Logger.Info("Setup")
				// setup.Setup(payload)
			default:
				errmsg := fmt.Sprintf("Unknown request type: %s", payload.Type)
				config.Logger.Warn(errmsg)
				return
			}
		}
	}
	// default to unauth
	if !auth.WsAuthCheck(conn) {
		unauthHandler(conn, r)
	}
}

// validate password and add to auth session map
func loginHandler(conn *websocket.Conn, msg []byte, payload structs.WsPayload) error {
	config.Logger.Info("Login")
	payload.Payload = structs.WsLoginPayload{}
	if err := json.Unmarshal(msg, &payload); err != nil {
		return fmt.Errorf("Error unmarshalling message: %v", err)
	}
	loginPayload, ok := payload.Payload.(structs.WsLoginPayload)
	if !ok {
		return fmt.Errorf("Error casting to LoginPayload")
	}
	isAuthenticated := auth.AuthenticateLogin(loginPayload.Password)
	if isAuthenticated {
		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 {
		config.Logger.Info("Login failed")
		return fmt.Errorf("Failed auth")
	}
	return nil
}


// broadcast the unauth payload
func unauthHandler(conn *websocket.Conn, r *http.Request) {
	config.Logger.Info("Sending unauth broadcast")
	blob := structs.UnauthBroadcast{
        Type:      "structure",
        AuthLevel: "unauthorized",
        Login: struct {
            Remainder int `json:"remainder"`
        }{
            Remainder: 0,
        },
    }
	resp, err := json.Marshal(blob)
	if err != nil {
		config.Logger.Error(fmt.Sprintf("Error unmarshalling message: %v", err))
		return
	}
	if err := conn.WriteMessage(websocket.TextMessage, resp); err != nil {
		config.Logger.Error(fmt.Sprintf("Error writing response: %v", err))
		return
	}
}

// client send:
// {
// 	"type": "verify",
// 	"id": "jsgeneratedid",
// 	"token<optional>": {
// 	  "id": "servergeneratedid",
// 	  "token": "encryptedtext"
// 	}
// }

// 1. we decrypt the token
// 2. we modify token['authorized'] to true
// 3. remove it from 'unauthorized' in system.json
// 4. hash and add to 'authozired' in system.json
// 5. encrypt that, and send it back to the user

// server respond:
// {
// 	"type": "activity",
// 	"response": "ack/nack",
// 	"error": "null/<some_error>",
// 	"id": "jsgeneratedid",
// 	"token": { (either new token or the token the user sent us)
// 	  "id": "relevant_token_id",
// 	  "token": "encrypted_text"
// 	}
// }

func supportHandler(msg []byte, payload structs.WsPayload, r *http.Request, conn *websocket.Conn) error {
	config.Logger.Info("Support")
	return nil
}