refactor: finished refactoring server-side code

1 files changed, 137 insertions, 109 deletions

diff --git a/internal/server/server.go b/internal/server/server.go
index 933ad29..4ac8969 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -2,6 +2,7 @@ package server
 
 import (
 	"context"
+	"crypto/ed25519"
 	"crypto/rand"
 	"crypto/tls"
 	_ "embed"
@@ -10,6 +11,7 @@ import (
 	"io"
 	"log"
 	"net"
+	"os"
 	"strconv"
 	"strings"
 	"sync"
@@ -48,44 +50,64 @@ func init() {
 	}
 
 	message := packet.ErrorMessage{Error: packet.PacketUnsupportedEncoding.Error()}
-	encoder, err := packet.NewMsgPackEncoder(&message)
-	assert.NoError(err, "constant packets should not error")
-	unsupportedEncodingErrorPacket = packet.NewPacket(encoder)
+	unsupportedEncodingErrorPacket = packet.NewPacket(packet.NewMsgPackEncoder(&message))
 
 	message = packet.ErrorMessage{Error: packet.PacketUnsupportedType.Error()}
-	encoder, err = packet.NewMsgPackEncoder(&message)
-	assert.NoError(err, "constant packets should not error")
-	unsupportedTypeErrorPacket = packet.NewPacket(encoder)
+	unsupportedTypeErrorPacket = packet.NewPacket(packet.NewMsgPackEncoder(&message))
 }
 
 type server struct {
-	node *snowflake.Node
-	port uint16
+	Node     *snowflake.Node
+	Port     uint16
+	sessions map[snowflake.ID]*Session
+	sessMu   sync.RWMutex
 }
 
 func NewServer(port uint16) server {
-	assert.Assert(nodeId <= snowflake.NodeMax, "maximum amount of servers reached %v", snowflake.NodeMax)
+	assert.Assert(nodeId <= snowflake.NodeMax, "maximum amount of servers reached")
 	node := snowflake.NewNode(nodeId)
 	nodeId++
 
 	return server{
-		node: node,
-		port: port,
+		Node:     node,
+		Port:     port,
+		sessions: map[snowflake.ID]*Session{},
 	}
 }
 
+func (s *server) AddSession(session *Session) {
+	s.sessMu.Lock()
+	defer s.sessMu.Unlock()
+	s.sessions[session.ID] = session
+}
+
+func (s *server) RemoveSession(id snowflake.ID) {
+	s.sessMu.Lock()
+	defer s.sessMu.Unlock()
+	delete(s.sessions, id)
+}
+
+func (s *server) Session(id snowflake.ID) (*Session, bool) {
+	s.sessMu.RLock()
+	defer s.sessMu.RUnlock()
+	session, ok := s.sessions[id]
+	return session, ok
+}
+
 func (s *server) ListenAndServe(ctx context.Context) {
-	listener, err := tls.Listen("tcp4", ":"+strconv.Itoa(int(s.port)), tlsConfig)
+	listener, err := tls.Listen("tcp4", ":"+strconv.Itoa(int(s.Port)), tlsConfig)
 	if err != nil {
 		log.Fatalf("error starting server: %s", err)
 	}
+
+	assert.AddFlush(listener)
 	defer listener.Close()
 	go func() {
 		<-ctx.Done()
 		listener.Close()
 	}()
 
-	log.Printf("started listening on port %v...\n", s.port)
+	log.Printf("started listening on port %v...\n", s.Port)
 	var wg sync.WaitGroup
 	for {
 		conn, err := listener.Accept()
@@ -97,46 +119,61 @@ func (s *server) ListenAndServe(ctx context.Context) {
 		}
 		wg.Add(1)
 		go func() {
-			handleConnection(ctx, conn)
+			handleConnection(ctx, conn, s)
 			wg.Done()
 		}()
 	}
-	log.Printf("stopped listening on port %v\n", s.port)
+	log.Printf("stopped listening on port %v\n", s.Port)
 
 	log.Println("waiting for all active connections to close...")
 	wg.Wait()
 	log.Println("server shutdown complete")
 }
 
-func handleConnection(ctx context.Context, conn net.Conn) {
+func handleConnection(ctx context.Context, conn net.Conn, server *server) {
 	addr, ok := conn.RemoteAddr().(*net.TCPAddr)
 	assert.Assert(ok, "getting tcp address should never fail as we are using tcp connections")
 
-	writeQueue := make(chan packet.Packet, 10)
-	session := newSession(addr, writeQueue)
-	nonce := session.Challenge()
+	log.Println(addr, "accepted")
 
+	nonce := [32]byte{}
+	_, err := rand.Read(nonce[:])
+	assert.NoError(err, "random should always produce a value")
+	pubKey, err := handleAuth(conn, nonce[:])
+	if err != nil {
+		log.Println(addr, err)
+		conn.Close()
+		log.Println(addr, "disconnected")
+		return
+	}
+
+	// TODO: replace this with DB query for id
+	id := server.Node.Generate()
+	session := newSession(server, addr, id, pubKey)
+	server.AddSession(session)
 	ctx = newContext(ctx, session)
 	framer := packet.NewFramer(ctx)
 
-	log.Println(addr, "accepted")
-
 	defer func() {
 		conn.Close()
+		close(framer.Out)
+		server.RemoveSession(session.ID)
+		close(session.WriteQueue)
 		log.Println(addr, "disconnected")
 	}()
 
 	go func() {
-		var mu sync.Mutex
 		for {
-			packet, ok := <-writeQueue
+			packet, ok := <-session.WriteQueue
 			if !ok {
 				break
 			}
-			mu.Lock()
-			packet.Into(conn)
-			mu.Unlock()
+			if _, err := packet.Into(conn); err != nil {
+				log.Println(addr, err)
+				break
+			}
 		}
+		session.WriteQueue = nil
 	}()
 
 	go func() {
@@ -146,91 +183,97 @@ func handleConnection(ctx context.Context, conn net.Conn) {
 				break
 			}
 			response := processPacket(ctx, request)
-			writeQueue <- response
+			if session.WriteQueue == nil {
+				break
+			}
+			session.WriteQueue <- response
 		}
 	}()
 
 	buffer := make([]byte, 512)
 	for {
+		conn.SetReadDeadline(time.Now().Add(time.Second))
 		n, err := conn.Read(buffer)
-		if err != nil {
+		deadlineExceeded := errors.Is(err, os.ErrDeadlineExceeded)
+		if err != nil && !deadlineExceeded {
 			if !errors.Is(err, io.EOF) {
 				log.Println(addr, "read error:", err)
 			}
 			break
 		}
 
+		if ctx.Err() != nil {
+			log.Println(addr, ctx.Err())
+			break
+		}
+
 		err = framer.Push(ctx, buffer[:n])
 		if err != nil {
-			// Wrap err and send to client then break
+			if ctx.Err() != nil {
+				log.Println(addr, ctx.Err())
+			} else {
+				// TODO: Wrap err and send to client then break
+			}
+			break
 		}
 	}
 }
 
-func _handleConnection(ctx context.Context, conn net.Conn) {
-	addr, ok := conn.RemoteAddr().(*net.TCPAddr)
-	assert.Assert(ok, "getting tcp address should be valid")
+func handleAuth(conn net.Conn, nonce []byte) (ed25519.PublicKey, error) {
+	conn.SetDeadline(time.Now().Add(time.Second * 5))
 
-	log.Println(addr, "accepted")
-	defer log.Println(addr, "disconnected")
-	defer conn.Close()
+	challengePacket := make([]byte, len(nonce)+1)
+	challengePacket[0] = packet.VERSION
+	copy(challengePacket[1:], nonce)
 
-	// ctx = newContext(ctx, addr)
-	// // TODO: consider adding timeout/deadline to ctx?
+	_, err := conn.Write(challengePacket)
+	if err != nil {
+		return nil, fmt.Errorf("error writing challenge: %w", err)
+	}
 
-	out, outErr := packet.RunFramer(ctx, conn)
-outer:
-	for {
-		select {
-		case packet, ok := <-out:
-			if !ok {
-				break outer
-			}
-			log.Printf("client %v: request packet: %v\n", conn.RemoteAddr().String(), packet)
-			responsePacket, err := handlePacket(packet)
-			log.Printf("client %v: response packet: %v\n", conn.RemoteAddr().String(), responsePacket)
-			if err != nil {
-				log.Printf("client %v: error processing request: %v\n", conn.RemoteAddr().String(), err)
-				break outer
-			}
-			_, err = responsePacket.Into(conn)
-			if err != nil {
-				log.Printf("client %v: error writing packet: %v\n", conn.RemoteAddr().String(), err)
-				break outer
-			}
+	challengeResponsePacket := make([]byte, ed25519.PublicKeySize+ed25519.SignatureSize+1)
+	bytesRead := 0
+	for bytesRead < len(challengeResponsePacket) {
+		n, err := conn.Read(challengeResponsePacket[bytesRead:])
+		if err != nil {
+			return nil, fmt.Errorf("error reading challenge response: %w", err)
+		}
+		bytesRead += n
+	}
 
-		case err := <-outErr:
-			if err == packet.PacketUnsupportedEncoding {
-				_, err := unsupportedEncodingErrorPacket.Into(conn)
-				log.Printf("client %v: error writing unsupported encoding packet: %v\n", conn.RemoteAddr().String(), err)
-			} else if err == packet.PacketUnsupportedType {
-				_, err := unsupportedTypeErrorPacket.Into(conn)
-				log.Printf("client %v: error writing unsupported type packet: %v\n", conn.RemoteAddr().String(), err)
-			} else if err != nil {
-				log.Printf("client %v: internal error: %v\n", conn.RemoteAddr().String(), err)
-			}
-			break outer
+	if challengeResponsePacket[0] != packet.VERSION {
+		return nil, fmt.Errorf("incompatible version: %v", challengeResponsePacket[0])
+	}
 
-		case <-ctx.Done():
-			log.Printf("client %v: %v\n", conn.RemoteAddr().String(), ctx.Err())
-			break outer
-		}
+	pubKey := ed25519.PublicKey(challengeResponsePacket[1 : 1+ed25519.PublicKeySize])
+	signature := ed25519.PrivateKey(challengeResponsePacket[1+ed25519.PublicKeySize:])
+
+	if ok := ed25519.Verify(pubKey, nonce, signature); !ok {
+		return nil, errors.New("signature verification failed")
 	}
+
+	conn.SetDeadline(time.Time{})
+	return pubKey, nil
 }
 
 type Session struct {
+	Server     *server
 	Addr       *net.TCPAddr
-	WriteQueue <-chan packet.Packet
+	WriteQueue chan packet.Packet
+	ID         snowflake.ID
+	PubKey     ed25519.PublicKey
 
 	mu         sync.Mutex
-	challenge   []byte
+	challenge  []byte
 	issuedTime time.Time
 }
 
-func newSession(addr *net.TCPAddr, writeQueue <-chan packet.Packet) *Session {
+func newSession(server *server, addr *net.TCPAddr, id snowflake.ID, pubKey ed25519.PublicKey) *Session {
 	session := &Session{
+		Server:     server,
 		Addr:       addr,
-		WriteQueue: writeQueue,
+		PubKey:     pubKey,
+		WriteQueue: make(chan packet.Packet, 10),
 		challenge:  make([]byte, 32), // Recommended nonce size
 	}
 	session.Challenge() // Make sure an initial nonce is generated
@@ -263,58 +306,43 @@ func FromContext(ctx context.Context) (*Session, bool) {
 
 // TODO: Move everything below this to somewhere else
 
-func handlePacket(pkt packet.Packet) (packet.Packet, error) {
+func processPacket(ctx context.Context, pkt packet.Packet) packet.Packet {
+	session, ok := FromContext(ctx)
+	assert.Assert(ok, "context in process packet should always have a session")
+
 	var response packet.TypedMessage
 	switch pkt.Type() {
-	case packet.PacketEko:
-		var request packet.EkoMessage
-		if err := pkt.DecodePayload(&request); err != nil {
-			return packet.Packet{}, fmt.Errorf("decode error: %v", err)
-		}
-
-		response = &packet.EkoMessage{Message: "Eko \"" + request.Message + "\""}
 	case packet.PacketSendMessage:
-		var request packet.SendMessageMessage
+		var request packet.SendMessage
 		if err := pkt.DecodePayload(&request); err != nil {
-			return packet.Packet{}, fmt.Errorf("decode error: %v", err)
+			log.Println("decode error:", err)
+			response = &packet.ErrorMessage{Error: "malformed payload"}
+			break
 		}
 
 		content := strings.TrimSpace(request.Content)
 		if content == "" {
-			response = &packet.ErrorMessage{Error: "content must not be blank"}
+			response = &packet.ErrorMessage{Error: "message content must not be blank"}
 			break
 		}
 
+		node := session.Server.Node
 		message := data.Message{
 			Id:          node.Generate(),
-			SenderId:    node.Generate(),
-			FrequencyId: node.Generate(),
-			NetworkId:   node.Generate(),
+			SenderId:    session.ID,
+			FrequencyId: node.Generate(), // TODO: replace with actual ID
+			NetworkId:   node.Generate(), // TODO: replace with actual ID
 			Contents:    content,
 		}
 		messages = append(messages, message)
 
-		response = &packet.EkoMessage{Message: "Eko OK"}
-	case packet.PacketGetMessages:
-		var request packet.GetMessagesMessage
-		if err := pkt.DecodePayload(&request); err != nil {
-			return packet.Packet{}, fmt.Errorf("decode error: %v", err)
-		}
-
-		response = &packet.MessagesMessage{Messages: messages}
+		response = packet.NewOkMessage()
 	default:
-		return packet.Packet{}, errors.New("TODO: not implemented yet")
+		response = &packet.ErrorMessage{Error: "use of unsupported packet type"}
 	}
 
-	assert.NotNil(response, "response must always be set")
-	encoder, err := packet.NewMsgPackEncoder(response)
-	if err != nil {
-		return packet.Packet{}, fmt.Errorf("encode error: %v", err)
-	}
-	return packet.NewPacket(encoder), nil
+	assert.NotNil(response, "response must always be assigned to")
+	return packet.NewPacket(packet.NewMsgPackEncoder(response))
 }
 
-var (
-	node     = snowflake.NewNode(1)
-	messages []data.Message
-)
+var messages []data.Message