goseg/ws/ws.go

package ws

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/websocket"
	"goseg/auth"
	"goseg/broadcast"
	"goseg/config"
	"goseg/structs"
	"net/http"
	"time"
)

var (
	upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
	}
)

// 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) {
	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)
	// 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 prelim structs.WsType
		var payload structs.WsPayload
		if err := json.Unmarshal(msg, &prelim); err != nil {
			fmt.Println("Error unmarshalling message:", err)
			continue
		}
		switch prelim.Payload.Type {
		case "login":
			if err = loginHandler(msg, payload); err != nil {
				config.Logger.Error(fmt.Sprintf("%v", err))
			}
		case "verify":
			if err = verifyHandler(msg, payload, r, conn); err != nil {
				config.Logger.Error(fmt.Sprintf("%v", err))
			}
		case "setup":
			config.Logger.Info("Setup")
			// setup.Setup(payload)
		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":
			config.Logger.Info("Support")
		case "broadcast":
			if err := broadcast.BroadcastToClients(); err != nil {
				errmsg := fmt.Sprintf("Unable to broadcast to peer(s): %v", err)
				config.Logger.Error(errmsg)
			}
		default:
			errmsg := fmt.Sprintf("Unknown request type: %s", prelim.Payload.Type)
			config.Logger.Warn(errmsg)
		}
	}
}

// validate password and add to auth session map
func loginHandler(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)
	}
	loginPayload, ok := payload.Payload.(structs.WsLoginPayload)
	if !ok {
		return fmt.Errorf("Error casting to LoginPayload")
	}
	isAuthenticated := auth.AuthenticateLogin(loginPayload.Password)
	if isAuthenticated {
		if err := auth.AddSession(payload.Token.ID, payload.Token.Token, now, true); err != nil {
			return fmt.Errorf("Unable to process login: %v", err)
		}
	} else {
		config.Logger.Info("Login failed")
	}
	return nil
}

// create tokens for those who have not auth
func verifyHandler(msg []byte, payload structs.WsPayload, r *http.Request, conn *websocket.Conn) error {
	config.Logger.Info("Verify")
	payload.Payload = structs.WsLoginPayload{}
	// if we can't unmarshal, assume no token
	if err := json.Unmarshal(msg, &payload); err != nil {
		resp, err := auth.CreateToken(conn, r, false)
		if err != nil {
			fmt.Errorf("Couldn't create token: %v", err)
		}
		respJson, err := json.Marshal(resp)
		if err != nil {
			return fmt.Errorf("Error marshalling token: %v", err)
		}
		if err := conn.WriteMessage(websocket.TextMessage, respJson); err != nil {
			return fmt.Errorf("Error writing response: %v", err)
		}
	}
	return nil
}