Delegate session resumption to SessionResumer

Extract full challenge-response and transport restoration into a
SessionResumer and delegate handling from Dialer and Server, removing
duplicated verification/send logic. Refresh handshake UpdatedAt on load
so expiry restarts, switch SessionManager index loading to sync.Once,
and add error handling for GetChatHistory queries.

Author: hossein1376

dial.go

@@ -1,17 +1,14 @@
 package kamune
 
 import (
-	"crypto/subtle"
 	"fmt"
 	"log/slog"
 	"net"
 	"runtime/debug"
 	"time"
 
 	"github.com/xtaci/kcp-go/v5"
-	"google.golang.org/protobuf/proto"
 
-	"github.com/kamune-org/kamune/internal/box/pb"
 	"github.com/kamune-org/kamune/pkg/attest"
 	"github.com/kamune-org/kamune/pkg/fingerprint"
 )
@@ -239,146 +236,22 @@ func (d *Dialer) attemptResumption(state *SessionState) (*Transport, error) {
 		d.conn = c
 	}
 
-	// Generate a challenge for the server to prove it has the shared secret
-	challenge := randomBytes(resumeChallengeSize)
-
-	// Create and send reconnect request
-	req := &pb.ReconnectRequest{
-		SessionId:        state.SessionID,
-		LastPhase:        state.Phase.ToProto(),
-		LastSendSequence: state.SendSequence,
-		LastRecvSequence: state.RecvSequence,
-		RemotePublicKey:  d.attester.PublicKey().Marshal(),
-		ResumeChallenge:  challenge,
-	}
-	if err := d.sendSignedMessage(req, RouteReconnect); err != nil {
-		return nil, fmt.Errorf("sending reconnect request: %w", err)
-	}
-
-	// Receive response
-	respPayload, err := d.conn.ReadBytes()
-	if err != nil {
-		return nil, fmt.Errorf("reading reconnect response: %w", err)
-	}
-
-	var respST pb.SignedTransport
-	if err := proto.Unmarshal(respPayload, &respST); err != nil {
-		return nil, fmt.Errorf("unmarshaling response transport: %w", err)
-	}
-
-	// Verify signature from the server
-	remoteKey, err := d.storage.algorithm.Identitfier().ParsePublicKey(
-		state.RemotePublicKey,
-	)
-	if err != nil {
-		return nil, fmt.Errorf("parsing remote public key: %w", err)
-	}
-
-	if !d.storage.algorithm.Identitfier().Verify(
-		remoteKey, respST.Data, respST.Signature,
-	) {
-		return nil, ErrInvalidSignature
-	}
-
-	var resp pb.ReconnectResponse
-	if err := proto.Unmarshal(respST.Data, &resp); err != nil {
-		return nil, fmt.Errorf("unmarshaling response: %w", err)
-	}
-
-	if !resp.Accepted {
-		return nil, fmt.Errorf("%w: %s", ErrResumptionFailed, resp.ErrorMessage)
-	}
-
-	// Verify the server's challenge response
+	// Delegate the entire challenge-response protocol to SessionResumer.
 	resumer := NewSessionResumer(
 		d.storage,
 		d.sessionManager,
 		d.attester,
 		d.resumptionConfig.MaxSessionAge,
 	)
-	expectedResponse, err := resumer.computeChallengeResponse(
-		challenge, state.SharedSecret,
-	)
-	if err != nil {
-		return nil, fmt.Errorf("compute expected challenge response: %w", err)
-	}
-	if subtle.ConstantTimeCompare(resp.ChallengeResponse, expectedResponse) != 1 {
-		return nil, ErrChallengeVerifyFailed
-	}
-
-	// Compute our response to the server's challenge
-	clientChallengeResponse, err := resumer.computeChallengeResponse(
-		resp.ServerChallenge, state.SharedSecret,
-	)
-	if err != nil {
-		return nil, fmt.Errorf("compute client challenge response: %w", err)
-	}
-
-	// Determine the sequence numbers to use
-	resumeSendSeq, resumeRecvSeq := resumer.reconcileSequences(
-		state.SendSequence,
-		state.RecvSequence,
-		resp.ServerRecvSequence,
-		resp.ServerSendSequence,
-	)
 
-	// Send verification
-	verify := &pb.ReconnectVerify{
-		ChallengeResponse: clientChallengeResponse,
-		Verified:          true,
-	}
-	if err := d.sendSignedMessage(verify, RouteReconnect); err != nil {
-		return nil, fmt.Errorf("sending verification: %w", err)
-	}
-
-	// Receive completion
-	completePayload, err := d.conn.ReadBytes()
+	transport, err := resumer.InitiateResumption(d.conn, state)
 	if err != nil {
-		return nil, fmt.Errorf("reading completion: %w", err)
-	}
-
-	var completeST pb.SignedTransport
-	if err := proto.Unmarshal(completePayload, &completeST); err != nil {
-		return nil, fmt.Errorf("unmarshaling completion transport: %w", err)
-	}
-
-	if !d.storage.algorithm.Identitfier().Verify(
-		remoteKey, completeST.Data, completeST.Signature,
-	) {
-		return nil, ErrInvalidSignature
-	}
-
-	var complete pb.ReconnectComplete
-	if err := proto.Unmarshal(completeST.Data, &complete); err != nil {
-		return nil, fmt.Errorf("unmarshaling completion: %w", err)
-	}
-
-	if !complete.Success {
-		return nil, fmt.Errorf(
-			"%w: %s", ErrResumptionFailed, complete.ErrorMessage,
-		)
-	}
-
-	// Use the agreed-upon sequence numbers
-	if complete.ResumeSendSequence > 0 {
-		resumeSendSeq = complete.ResumeSendSequence
-	}
-	if complete.ResumeRecvSequence > 0 {
-		resumeRecvSeq = complete.ResumeRecvSequence
-	}
-
-	// Restore the transport
-	transport, err := resumer.restoreTransport(
-		d.conn, state, resumeSendSeq, resumeRecvSeq,
-	)
-	if err != nil {
-		return nil, fmt.Errorf("restoring transport: %w", err)
+		return nil, fmt.Errorf("initiating resumption: %w", err)
 	}
 
 	// Update session state
 	if d.resumptionConfig.PersistSessions {
-		err := SaveSessionForResumption(transport, d.sessionManager)
-		if err != nil {
+		if err := SaveSessionForResumption(transport, d.sessionManager); err != nil {
 			slog.Warn(
 				"failed to update session after resumption",
 				slog.Any("error", err),
@@ -389,32 +262,6 @@ func (d *Dialer) attemptResumption(state *SessionState) (*Transport, error) {
 	return transport, nil
 }
 
-func (d *Dialer) sendSignedMessage(msg Transferable, route Route) error {
-	data, err := proto.Marshal(msg)
-	if err != nil {
-		return fmt.Errorf("marshaling message: %w", err)
-	}
-
-	sig, err := d.attester.Sign(data)
-	if err != nil {
-		return fmt.Errorf("signing message: %w", err)
-	}
-
-	st := &pb.SignedTransport{
-		Data:      data,
-		Signature: sig,
-		Padding:   padding(maxPadding),
-		Route:     route.ToProto(),
-	}
-
-	payload, err := proto.Marshal(st)
-	if err != nil {
-		return fmt.Errorf("marshaling transport: %w", err)
-	}
-
-	return d.conn.WriteBytes(payload)
-}
-
 // PublicKey returns the dialer's public key.
 func (d *Dialer) PublicKey() PublicKey {
 	return d.attester.PublicKey()

server.go

@@ -1,7 +1,6 @@
 package kamune
 
 import (
-	"crypto/subtle"
 	"errors"
 	"fmt"
 	"log/slog"
@@ -213,13 +212,14 @@ func (s *Server) handleReconnection(cn Conn, st *pb.SignedTransport) error {
 		return s.sendReconnectReject(cn, "session resumption is disabled")
 	}
 
-	// Parse the reconnect request from the signed transport
+	// Parse the reconnect request from the signed transport.
 	var req pb.ReconnectRequest
 	if err := proto.Unmarshal(st.Data, &req); err != nil {
 		return fmt.Errorf("unmarshaling reconnect request: %w", err)
 	}
 
-	// Verify the signature using the claimed public key
+	// Verify the signature using the claimed public key before handing off
+	// to the resumer, which trusts the request is authentic.
 	remoteKey, err := s.algorithm.Identitfier().ParsePublicKey(
 		req.RemotePublicKey,
 	)
@@ -235,134 +235,24 @@ func (s *Server) handleReconnection(cn Conn, st *pb.SignedTransport) error {
 		slog.String("session_id", req.SessionId),
 	)
 
-	// Create session resumer and handle the resumption
+	// Delegate the entire challenge-response protocol to SessionResumer.
 	resumer := NewSessionResumer(
 		s.storage,
 		s.sessionManager,
 		s.attester,
 		s.resumptionConfig.MaxSessionAge,
 	)
 
-	// Look up the session
-	state, err := s.sessionManager.LoadSessionByPublicKey(req.RemotePublicKey)
+	t, err := resumer.HandleResumption(cn, &req)
 	if err != nil {
-		slog.Warn("session not found for reconnection",
-			slog.Any("error", err),
-		)
-		return s.sendReconnectReject(cn, "session not found")
-	}
-
-	// Verify session ID matches
-	if state.SessionID != req.SessionId {
-		return s.sendReconnectReject(cn, "session ID mismatch")
-	}
-
-	// Verify the session is in a resumable state
-	if state.Phase != PhaseEstablished {
-		return s.sendReconnectReject(cn, "session not established")
-	}
-
-	// Verify the session is not too old
-	// Note: We'd need to track creation time in SessionState for this
-	// For now, just check that we have a valid shared secret
-	if len(state.SharedSecret) == 0 {
-		return s.sendReconnectReject(cn, "session state invalid")
-	}
-
-	// Generate server challenge
-	serverChallenge := randomBytes(resumeChallengeSize)
-
-	// Compute response to client's challenge using HMAC
-	challengeResponse, err := resumer.computeChallengeResponse(
-		req.ResumeChallenge, state.SharedSecret,
-	)
-	if err != nil {
-		return fmt.Errorf("compute challenge response: %w", err)
-	}
-
-	// Send accept response
-	resp := &pb.ReconnectResponse{
-		Accepted:           true,
-		ResumeFromPhase:    state.Phase.ToProto(),
-		ChallengeResponse:  challengeResponse,
-		ServerChallenge:    serverChallenge,
-		ServerSendSequence: state.SendSequence,
-		ServerRecvSequence: state.RecvSequence,
-	}
-
-	if err := s.sendSignedMessage(cn, resp, RouteReconnect); err != nil {
-		return fmt.Errorf("sending accept response: %w", err)
-	}
-
-	// Receive client verification
-	verifyPayload, err := cn.ReadBytes()
-	if err != nil {
-		return fmt.Errorf("reading verification: %w", err)
-	}
-
-	var verifyST pb.SignedTransport
-	if err := proto.Unmarshal(verifyPayload, &verifyST); err != nil {
-		return fmt.Errorf("unmarshaling verification transport: %w", err)
-	}
-
-	// Verify signature
-	if !s.algorithm.Identitfier().Verify(
-		remoteKey, verifyST.Data, verifyST.Signature,
-	) {
-		return fmt.Errorf("verification signature invalid")
-	}
-
-	var verify pb.ReconnectVerify
-	if err := proto.Unmarshal(verifyST.Data, &verify); err != nil {
-		return fmt.Errorf("unmarshaling verification: %w", err)
-	}
-
-	// Verify client's response to our challenge
-	expectedClientResponse, err := resumer.computeChallengeResponse(
-		serverChallenge, state.SharedSecret,
-	)
-	if err != nil {
-		return fmt.Errorf("compute client expected response: %w", err)
-	}
-	if len(verify.ChallengeResponse) == 0 ||
-		subtle.ConstantTimeCompare(verify.ChallengeResponse, expectedClientResponse) != 1 {
-		complete := &pb.ReconnectComplete{
-			Success:      false,
-			ErrorMessage: "challenge verification failed",
-		}
-		_ = s.sendSignedMessage(cn, complete, RouteReconnect)
-		return ErrChallengeVerifyFailed
-	}
-
-	// Determine sequence numbers to resume from
-	resumeSendSeq, resumeRecvSeq := resumer.reconcileSequences(
-		state.SendSequence, state.RecvSequence,
-		req.LastRecvSequence, req.LastSendSequence,
-	)
-
-	// Send completion
-	complete := &pb.ReconnectComplete{
-		Success:            true,
-		ResumeSendSequence: resumeSendSeq,
-		ResumeRecvSequence: resumeRecvSeq,
-	}
-	if err := s.sendSignedMessage(cn, complete, RouteReconnect); err != nil {
-		return fmt.Errorf("sending completion: %w", err)
-	}
-
-	// Restore the transport
-	t, err := resumer.restoreTransport(cn, state, resumeSendSeq, resumeRecvSeq)
-	if err != nil {
-		return fmt.Errorf("restoring transport: %w", err)
+		return fmt.Errorf("handling resumption: %w", err)
 	}
 
 	slog.Info("session resumed",
 		slog.String("session_id", t.SessionID()),
-		slog.Uint64("send_seq", resumeSendSeq),
-		slog.Uint64("recv_seq", resumeRecvSeq),
 	)
 
-	// Update session state
+	// Update session state.
 	if s.resumptionConfig.PersistSessions {
 		if err := SaveSessionForResumption(t, s.sessionManager); err != nil {
 			slog.Warn("failed to update session after resumption",