Implemented server-side authentication in accordance to the protocol

5 files changed, 139 insertions, 121 deletions

diff --git a/internal/server/api/api.go b/internal/server/api/api.go
index 068e074..2640433 100644
--- a/internal/server/api/api.go
+++ b/internal/server/api/api.go
@@ -14,6 +14,7 @@ import (
 
 	"github.com/kyren223/eko/internal/data"
 	"github.com/kyren223/eko/internal/packet"
+	"github.com/kyren223/eko/internal/server/ctxkeys"
 	"github.com/kyren223/eko/internal/server/session"
 	"github.com/kyren223/eko/pkg/assert"
 	"github.com/kyren223/eko/pkg/snowflake"
@@ -393,12 +394,13 @@ func CreateNetwork(ctx context.Context, sess *session.Session, request *packet.C
 	}
 }
 
-func GetNetworksInfo(ctx context.Context, sess *session.Session) (packet.Payload, error) {
+func GetNetworksInfo(ctx context.Context, sess *session.Session) packet.Payload {
 	var fullNetworks []packet.FullNetwork
 
 	tx, err := db.BeginTx(ctx, nil)
 	if err != nil {
-		return nil, err
+		slog.ErrorContext(ctx, "database error", "error", err)
+		return &ErrInternalError
 	}
 	defer func() { _ = tx.Rollback() }()
 
@@ -407,18 +409,21 @@ func GetNetworksInfo(ctx context.Context, sess *session.Session) (packet.Payload
 
 	networks, err := qtx.GetUserNetworks(ctx, sess.ID())
 	if err != nil {
-		return nil, err
+		slog.ErrorContext(ctx, "database error", "error", err)
+		return &ErrInternalError
 	}
 
 	for _, network := range networks {
 		frequencies, err := qtx.GetNetworkFrequencies(ctx, network.ID)
 		if err != nil {
-			return nil, err
+			slog.ErrorContext(ctx, "database error", "error", err)
+			return &ErrInternalError
 		}
 
 		membersAndUsers, err := qtx.GetNetworkMembers(ctx, network.ID)
 		if err != nil {
-			return nil, err
+			slog.ErrorContext(ctx, "database error", "error", err)
+			return &ErrInternalError
 		}
 		members, users := SplitMembersAndUsers(membersAndUsers)
 
@@ -432,14 +437,15 @@ func GetNetworksInfo(ctx context.Context, sess *session.Session) (packet.Payload
 
 	err = tx.Commit()
 	if err != nil {
-		return nil, err
+		slog.ErrorContext(ctx, "database error", "error", err)
+		return &ErrInternalError
 	}
 
 	return &packet.NetworksInfo{
 		Networks:        fullNetworks,
 		RemovedNetworks: nil,
 		Partial:         false,
-	}, nil
+	}
 }
 
 func CreateFrequency(ctx context.Context, sess *session.Session, request *packet.CreateFrequency) packet.Payload {
@@ -1565,3 +1571,99 @@ func GetUsers(ctx context.Context, sess *session.Session, request *packet.GetUse
 		Users: users,
 	}
 }
+
+func GetNonce(ctx context.Context, sess *session.Session, request *packet.GetNonce) packet.Payload {
+	return &packet.NonceInfo{
+		Nonce: sess.Challenge(),
+	}
+}
+
+func Authenticate(ctx context.Context, sess *session.Session, request *packet.Authenticate) packet.Payload {
+	if sess.IsAuthenticated() {
+		return &packet.Error{Error: "already authenticated"}
+	}
+
+	if len(request.PubKey) != ed25519.PublicKeySize {
+		return &packet.Error{Error: fmt.Sprintf(
+			"public key must be exactly %v bytes", ed25519.PublicKeySize,
+		)}
+	}
+
+	if len(request.Signature) != ed25519.SignatureSize {
+		return &packet.Error{Error: fmt.Sprintf(
+			"signature must be exactly %v bytes", ed25519.SignatureSize,
+		)}
+	}
+
+	// IMPORTANT
+	if ok := ed25519.Verify(request.PubKey, sess.Challenge(), request.Signature); !ok {
+		return &packet.Error{Error: "signature verification failed"}
+	}
+
+	// Authenticated from here on out
+
+	queries := data.New(db)
+
+	user, err := queries.GetUserByPublicKey(ctx, request.PubKey)
+	if err == sql.ErrNoRows {
+		id := sess.Manager().Node().Generate()
+		user, err = queries.CreateUser(ctx, data.CreateUserParams{
+			ID:        id,
+			Name:      "User" + strconv.FormatInt(id.Time()%1000, 10),
+			PublicKey: request.PubKey,
+		})
+	}
+	if err != nil {
+		slog.ErrorContext(ctx, "database error", "error", err)
+		return &ErrInternalError
+	}
+
+	if user.IsDeleted {
+		return &packet.Error{Error: "public key is already taken by a deleted user"}
+	}
+
+	sess.Manager().AddSession(sess, user.ID, request.PubKey)
+
+	// NOTE: as per the protocol, this must be the first message after auth
+	payload := &packet.UsersInfo{Users: []data.User{user}}
+	ok := sess.Write(ctx, WrapPayload(payload))
+	if !ok {
+		// Timeout, send at least this payload, client can request the rest
+		return payload
+	}
+
+	ok = sendInitialAuthPackets(ctx, sess) // Send rest of packets
+	if !ok {
+		// Timeout, notify the client at least
+		return &packet.Error{Error: "timeout: not all initial auth packets were sent"}
+	}
+
+	return nil // manually writing requests to control order
+}
+
+func sendInitialAuthPackets(ctx context.Context, sess *session.Session) bool {
+	payloads := []packet.Payload{}
+
+	payloads = append(payloads, GetUserData(ctx, sess, &packet.GetUserData{}))
+	payloads = append(payloads, GetTrustedUsers(ctx, sess))
+	payloads = append(payloads, GetBlockedUsers(ctx, sess))
+	payloads = append(payloads, GetBlockedUsers(ctx, sess))
+	payloads = append(payloads, GetNetworksInfo(ctx, sess))
+	payloads = append(payloads, GetNotifications(ctx, sess))
+
+	success := true
+	for _, payload := range payloads {
+		if payload == &ErrInternalError {
+			success = false
+			continue
+		}
+		slog.InfoContext(ctx, "sending initial auth payload", ctxkeys.Payload.String(), payload, ctxkeys.PayloadType.String(), payload.Type())
+		ok := sess.Write(ctx, WrapPayload(payload))
+		if !ok {
+			success = false
+			continue
+		}
+	}
+
+	return success
+}
diff --git a/internal/server/api/helpers.go b/internal/server/api/helpers.go
index 4327a49..2ed0694 100644
--- a/internal/server/api/helpers.go
+++ b/internal/server/api/helpers.go
@@ -79,8 +79,7 @@ func NetworkPropagateWithFilter(
 		context, cancel := context.WithTimeout(context.Background(), timeout)
 		go func() {
 			defer cancel()
-			pkt := packet.NewPacket(packet.NewMsgPackEncoder(payload))
-			if ok := session.Write(context, pkt); !ok {
+			if ok := session.Write(context, WrapPayload(payload)); !ok {
 				log.Println(sess.Addr(), "propagation to", session.Addr(), "failed")
 			}
 		}()
@@ -122,8 +121,7 @@ func UserPropagate(
 	context, cancel := context.WithTimeout(context.Background(), timeout)
 	go func() {
 		defer cancel()
-		pkt := packet.NewPacket(packet.NewMsgPackEncoder(payload))
-		if ok := session.Write(context, pkt); !ok {
+		if ok := session.Write(context, WrapPayload(payload)); !ok {
 			log.Println(sess.Addr(), "propagation to", session.Addr(), "failed")
 		}
 	}()
@@ -191,3 +189,7 @@ func getNotifications(ctx context.Context, userId snowflake.ID) (packet.Notifica
 	}
 	return items, nil
 }
+
+func WrapPayload(payload packet.Payload) packet.Packet {
+	return packet.NewPacket(packet.NewMsgPackEncoder(payload))
+}
diff --git a/internal/server/ctxkeys/ctxkeys.go b/internal/server/ctxkeys/ctxkeys.go
index 827e6ea..e832e78 100644
--- a/internal/server/ctxkeys/ctxkeys.go
+++ b/internal/server/ctxkeys/ctxkeys.go
@@ -14,8 +14,8 @@ const (
 	IpAddr
 	Evicted
 	EvictedBy
-	Request
-	RequestType
+	Payload
+	PayloadType
 
 	KeyMax
 )
@@ -25,8 +25,8 @@ var keyNames = map[key]string{
 	IpAddr:      "ip_addr",
 	Evicted:     "evicted",
 	EvictedBy:   "evicted_by",
-	Request:     "request",
-	RequestType: "request_type",
+	Payload:     "payload",
+	PayloadType: "payload_type",
 }
 
 func (k key) String() string {
diff --git a/internal/server/server.go b/internal/server/server.go
index cdb9313..41eab92 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -3,10 +3,8 @@ package server
 import (
 	"context"
 	"crypto/ed25519"
-	"crypto/rand"
 	"crypto/tls"
 	"errors"
-	"fmt"
 	"io"
 	"log"
 	"log/slog"
@@ -104,8 +102,7 @@ func EvictSession(sess *session.Session) {
 	payload := &packet.Error{
 		Error: "new connection from another location, closing this one",
 	}
-	pkt := packet.NewPacket(packet.NewMsgPackEncoder(payload))
-	sess.Write(ctx, pkt)
+	sess.Write(ctx, api.WrapPayload(payload))
 
 	sess.Close()
 }
@@ -240,9 +237,7 @@ func (server *server) handleConnection(conn net.Conn) {
 		localCtx := context.WithoutCancel(ctx)
 
 		for request := range framer.Out {
-			response := processPacket(localCtx, sess, request)
-			ok = sess.Write(localCtx, response)
-			assert.Assert(ok, "context is never done and write will panic")
+			processPacket(localCtx, sess, request)
 		}
 	}()
 
@@ -269,10 +264,8 @@ func (server *server) handleConnection(conn net.Conn) {
 			break
 		}
 		if err != nil {
-			payload := packet.Error{Error: err.Error()}
-			pkt := packet.NewPacket(packet.NewMsgPackEncoder(&payload))
 			writerWg.Add(1)
-			sess.Write(ctx, pkt)
+			sess.Write(ctx, api.WrapPayload(&packet.Error{Error: err.Error()}))
 			writerWg.Done()
 			slog.WarnContext(ctx, "received malformed packet", "error", err)
 			break
@@ -283,45 +276,7 @@ func (server *server) handleConnection(conn net.Conn) {
 	<-done
 }
 
-func handleAuth(conn net.Conn) (ed25519.PublicKey, error) {
-	nonce := [32]byte{}
-	_, err := rand.Read(nonce[:])
-	assert.NoError(err, "random should always produce a value")
-
-	challengePacket := make([]byte, len(nonce)+1)
-	challengePacket[0] = packet.VERSION
-	copy(challengePacket[1:], nonce[:])
-
-	_, err = conn.Write(challengePacket)
-	if err != nil {
-		return nil, fmt.Errorf("error writing challenge: %w", err)
-	}
-
-	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
-	}
-
-	if challengeResponsePacket[0] != packet.VERSION {
-		return nil, fmt.Errorf("incompatible version: %v", challengeResponsePacket[0])
-	}
-
-	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")
-	}
-
-	return pubKey, nil
-}
-
-func processPacket(ctx context.Context, sess *session.Session, pkt packet.Packet) packet.Packet {
+func processPacket(ctx context.Context, sess *session.Session, pkt packet.Packet) {
 	var response packet.Payload
 
 	request, err := pkt.DecodedPayload()
@@ -331,8 +286,11 @@ func processPacket(ctx context.Context, sess *session.Session, pkt packet.Packet
 		response = processRequest(ctx, sess, request)
 	}
 
-	assert.NotNil(response, "response must always be assigned to")
-	return packet.NewPacket(packet.NewMsgPackEncoder(response))
+	// Nil is ok if responses were handled manually using sess.Write()
+	if response != nil {
+		ok := sess.Write(ctx, api.WrapPayload(response))
+		assert.Assert(ok, "context is never done and write will panic if queue is closed")
+	}
 }
 
 func processRequest(ctx context.Context, sess *session.Session, request packet.Payload) packet.Payload {
@@ -371,16 +329,14 @@ func processRequest(ctx context.Context, sess *session.Session, request packet.P
 	}
 
 	switch request := request.(type) {
-	// TODO:
-	// case *packet.GetNonce:
-	// response = timeout(5*time.Millisecond, api.SetUserData, ctx, sess, request)
 
-	// TODO:
-	// case *packet.Authenticate:
-	// response = timeout(5*time.Millisecond, api.SetUserData, ctx, sess, request)
+	case *packet.GetNonce:
+		response = timeout(5*time.Millisecond, api.GetNonce, ctx, sess, request)
+
+	case *packet.Authenticate:
+		response = timeout(5*time.Millisecond, api.Authenticate, ctx, sess, request)
 
 	default:
-		_ = request // FIXME: remove
 		response = &packet.Error{Error: "use of disallowed packet type for request"}
 	}
 
@@ -485,50 +441,5 @@ func sendTosInfo(ctx context.Context, sess *session.Session) bool {
 		PrivacyPolicy: privacy,
 		Date:          date,
 	}
-	pkt := packet.NewPacket(packet.NewMsgPackEncoder(payload))
-	return sess.Write(ctx, pkt)
-}
-
-func (server *server) sendInitialAuthPackets(ctx context.Context, sess *session.Session) bool {
-	payload := api.GetUserData(ctx, sess, &packet.GetUserData{})
-	if payload == &api.ErrInternalError {
-		return false
-	}
-	log.Println(sess.Addr(), "sending", payload.Type(), "payload:", payload)
-	pkt := packet.NewPacket(packet.NewMsgPackEncoder(payload))
-	sess.Write(ctx, pkt)
-
-	payload = api.GetTrustedUsers(ctx, sess)
-	if payload == &api.ErrInternalError {
-		return false
-	}
-	log.Println(sess.Addr(), "sending", payload.Type(), "payload:", payload)
-	pkt = packet.NewPacket(packet.NewMsgPackEncoder(payload))
-	sess.Write(ctx, pkt)
-
-	payload = api.GetBlockedUsers(ctx, sess)
-	if payload == &api.ErrInternalError {
-		return false
-	}
-	log.Println(sess.Addr(), "sending", payload.Type(), "payload:", payload)
-	pkt = packet.NewPacket(packet.NewMsgPackEncoder(payload))
-	sess.Write(ctx, pkt)
-
-	payload, err := api.GetNetworksInfo(ctx, sess)
-	if err != nil {
-		return false
-	}
-	log.Println(sess.Addr(), "sending", payload.Type(), "payload:", payload)
-	pkt = packet.NewPacket(packet.NewMsgPackEncoder(payload))
-	sess.Write(ctx, pkt)
-
-	payload = api.GetNotifications(ctx, sess)
-	if payload == &api.ErrInternalError {
-		return false
-	}
-	log.Println(sess.Addr(), "sending", payload.Type(), "payload:", payload)
-	pkt = packet.NewPacket(packet.NewMsgPackEncoder(payload))
-	sess.Write(ctx, pkt)
-
-	return true
+	return sess.Write(ctx, api.WrapPayload(payload))
 }
diff --git a/internal/server/session/session.go b/internal/server/session/session.go
index 013ff71..b34d43a 100644
--- a/internal/server/session/session.go
+++ b/internal/server/session/session.go
@@ -13,7 +13,10 @@ import (
 	"github.com/kyren223/eko/pkg/snowflake"
 )
 
-const WriteQueueSize = 10
+const (
+	WriteQueueSize = 10
+	NonceSize      = 32
+)
 
 type SessionManager interface {
 	AddSession(session *Session, userId snowflake.ID, pubKey ed25519.PublicKey)
@@ -58,7 +61,7 @@ func NewSession(
 		writerWg:      writerWg,
 		writeMu:       sync.RWMutex{},
 		issuedTime:    time.Time{},
-		challenge:     make([]byte, 32),
+		challenge:     make([]byte, NonceSize),
 		pubKey:        ed25519.PublicKey{},
 		id:            snowflake.InvalidID,
 		mu:            sync.Mutex{},