Implemented proper connection and retry mechanisms

cmd/client/main.go

@@ -8,7 +8,7 @@ import (
 )
 
 func main() {
-	logFile, err := os.OpenFile("logs/client.log", os.O_APPEND|os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
+	logFile, err := os.OpenFile("client.log", os.O_APPEND|os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
 	if err != nil {
 		log.Fatalln(err)
 	}

cmd/server/main.go

@@ -19,7 +19,7 @@ func main() {
 	stdout := flag.Bool("stdout", false, "enable logging to stdout")
 	flag.Parse()
 
-	logFile, err := os.OpenFile("logs/server.log", os.O_APPEND|os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
+	logFile, err := os.OpenFile("server.log", os.O_APPEND|os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
 	if err != nil {
 		log.Fatalln(err)
 	}

internal/client/client.go

@@ -1,10 +1,13 @@
 package client
 
 import (
+	"io"
 	"log"
+	"os"
 	"reflect"
 
 	tea "github.com/charmbracelet/bubbletea"
+	"github.com/davecgh/go-spew/spew"
 	"github.com/muesli/termenv"
 
 	"github.com/kyren223/eko/internal/client/config"
@@ -23,6 +26,15 @@ func (c BubbleTeaCloser) Close() error {
 }
 
 func Run() {
+	var dump *os.File
+	if ui.DEBUG {
+		var err error
+		dump, err = os.OpenFile("messages.log", os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o644)
+		if err != nil {
+			os.Exit(1)
+		}
+	}
+
 	log.Println("client started")
 
 	err := config.Load()
@@ -36,7 +48,7 @@ func Run() {
 	// It only changes the current pane so new terminal panes/windows are not affected.
 	termenv.DefaultOutput().SetBackgroundColor(termenv.RGBColor(ui.BackgroundColor))
 
-	program := tea.NewProgram(initialModel(), tea.WithAltScreen())
+	program := tea.NewProgram(initialModel(dump), tea.WithAltScreen())
 	assert.AddFlush(BubbleTeaCloser{program})
 	ui.Program = program
 
@@ -46,11 +58,15 @@ func Run() {
 }
 
 type model struct {
+	dump  io.Writer
 	model tea.Model
 }
 
-func initialModel() model {
-	return model{auth.New()}
+func initialModel(dump io.WriteCloser) model {
+	return model{
+		dump:  dump,
+		model: auth.New(),
+	}
 }
 
 func (m model) Init() tea.Cmd {
@@ -62,10 +78,14 @@ func (m model) View() string {
 }
 
 func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	if m.dump != nil {
+		spew.Fdump(m.dump, msg)
+	}
+
 	switch msg := msg.(type) {
 	case tea.KeyMsg:
 		if msg.Type == tea.KeyCtrlC {
-			return m, tea.Quit
+			return m, func() tea.Msg { return ui.QuitMsg{} }
 		}
 
 	case tea.WindowSizeMsg:
@@ -80,5 +100,10 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	}
 	var cmd tea.Cmd
 	m.model, cmd = m.model.Update(msg)
+
+	if _, ok := msg.(ui.QuitMsg); ok {
+		return m, tea.Quit
+	}
+
 	return m, cmd
 }

internal/client/gateway/gateway.go

@@ -6,8 +6,8 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	_ "embed"
+	"encoding/binary"
 	"errors"
-	"io"
 	"log"
 	"net"
 	"sync"
@@ -18,6 +18,7 @@ import (
 	"github.com/kyren223/eko/internal/client/ui"
 	"github.com/kyren223/eko/internal/packet"
 	"github.com/kyren223/eko/pkg/assert"
+	"github.com/kyren223/eko/pkg/snowflake"
 )
 
 //go:embed server.crt
@@ -32,10 +33,11 @@ var (
 	framer  packet.PacketFramer
 	conn    net.Conn
 	writeMu sync.Mutex
+	closed  = false
 )
 
 type (
-	ConnectionEstablished struct{}
+	ConnectionEstablished snowflake.ID
 	ConnectionFailed      error
 	ConnectionLost        error
 	ConnectionClosed      struct{}
@@ -57,17 +59,19 @@ func Connect(privKey ed25519.PrivateKey, timeout time.Duration) tea.Cmd {
 	return func() tea.Msg {
 		ctx, cancel := context.WithTimeout(context.Background(), timeout)
 		defer cancel()
-		err := connect(ctx, privKey)
+		id, err := connect(ctx, privKey)
 		if err != nil {
 			return ConnectionFailed(err)
 		}
-		return ConnectionEstablished{}
+		return ConnectionEstablished(id)
 	}
 }
 
-func connect(ctx context.Context, privKey ed25519.PrivateKey) error {
+func connect(ctx context.Context, privKey ed25519.PrivateKey) (snowflake.ID, error) {
 	assert.Assert(conn == nil, "cannot connect, connection is active")
+	closed = false
 
+	var id snowflake.ID
 	connChan := make(chan net.Conn, 1)
 	errChan := make(chan error, 1)
 	go func() {
@@ -79,7 +83,7 @@ func connect(ctx context.Context, privKey ed25519.PrivateKey) error {
 		}
 		log.Println("established connection with server")
 
-		if err := handleAuth(ctx, privKey); err != nil {
+		if id, err = handleAuth(ctx, connection, privKey); err != nil {
 			errChan <- err
 			return
 		}
@@ -91,18 +95,18 @@ func connect(ctx context.Context, privKey ed25519.PrivateKey) error {
 	case connection := <-connChan:
 		conn = connection
 	case err := <-errChan:
-		return err
+		return 0, err
 	case <-ctx.Done():
-		return ctx.Err()
+		return 0, ctx.Err()
 	}
 
 	go readForever()
 	go handlePacketStream()
 
-	return nil
+	return id, nil
 }
 
-func handleAuth(ctx context.Context, privKey ed25519.PrivateKey) error {
+func handleAuth(ctx context.Context, conn net.Conn, privKey ed25519.PrivateKey) (snowflake.ID, error) {
 	const nonceSize = 32
 	challengeRequest := make([]byte, 1+nonceSize)
 
@@ -118,7 +122,7 @@ func handleAuth(ctx context.Context, privKey ed25519.PrivateKey) error {
 	for bytesRead < 1+nonceSize {
 		n, err := conn.Read(challengeRequest[bytesRead:])
 		if err != nil {
-			return err
+			return 0, err
 		}
 		bytesRead += n
 	}
@@ -134,10 +138,21 @@ func handleAuth(ctx context.Context, privKey ed25519.PrivateKey) error {
 
 	_, err = conn.Write(challengeResponse)
 	if err != nil {
-		return err
+		return 0, err
 	}
 
-	return nil
+	var idBytes [8]byte
+	bytesRead = 0
+	for bytesRead < 8 {
+		n, err := conn.Read(idBytes[:])
+		if err != nil {
+			return 0, err
+		}
+		bytesRead += n
+	}
+	id := snowflake.ID(binary.BigEndian.Uint64(idBytes[:]))
+
+	return id, nil
 }
 
 func readForever() {
@@ -145,9 +160,6 @@ func readForever() {
 	for conn != nil {
 		n, err := conn.Read(buffer)
 		if err != nil {
-			if errors.Is(err, io.EOF) {
-				err = nil
-			}
 			onDisconnect(err)
 			return
 		}
@@ -196,6 +208,7 @@ func handlePacketStream() {
 func Disconnect() {
 	assert.Assert(conn != nil, "cannot disconnect, connection is inactive")
 	conn.Close()
+	closed = true
 }
 
 func onDisconnect(err error) {
@@ -208,12 +221,12 @@ func onDisconnect(err error) {
 	}
 	asyncResponses = nil
 	responsesMu.Unlock()
-	if err != nil {
-		log.Println("connection lost:", err)
-		ui.Program.Send(ConnectionLost(err))
-	} else {
+	if closed {
 		log.Println("connection closed")
 		ui.Program.Send(ConnectionClosed{})
+	} else {
+		log.Println("connection lost:", err)
+		ui.Program.Send(ConnectionLost(err))
 	}
 }
 

internal/client/ui/core/core.go

@@ -2,37 +2,51 @@ package core
 
 import (
 	"crypto/ed25519"
+	"fmt"
 	"time"
 
+	"github.com/charmbracelet/bubbles/spinner"
+	"github.com/charmbracelet/bubbles/timer"
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/kyren223/eko/internal/client/gateway"
+	"github.com/kyren223/eko/internal/client/ui"
 	"github.com/kyren223/eko/internal/client/ui/loadscreen"
+	"github.com/kyren223/eko/pkg/snowflake"
 )
 
 var (
 	connectingToServer = "Connecting to server.."
 	connectionFailed   = "Connection failed - retrying in %d sec..."
+	connectionTimeout  = 5 * time.Second
+	initialTimeout     = 3750 * time.Millisecond
+	timerInterval      = 50 * time.Millisecond
 )
 
 type Model struct {
-	privKey ed25519.PrivateKey
 	name    string
+	privKey ed25519.PrivateKey
 
 	loading   loadscreen.Model
+	timeout   time.Duration
+	timer     timer.Model
 	connected bool
+
+	id snowflake.ID
 }
 
 func New(privKey ed25519.PrivateKey, name string) Model {
 	return Model{
-		privKey:   privKey,
 		name:      name,
+		privKey:   privKey,
 		loading:   loadscreen.New(connectingToServer),
+		timeout:   initialTimeout,
+		timer:     newTimer(initialTimeout),
 		connected: false,
 	}
 }
 
 func (m Model) Init() tea.Cmd {
-	return tea.Batch(gateway.Connect(m.privKey, 5*time.Second), m.loading.Init())
+	return tea.Batch(gateway.Connect(m.privKey, connectionTimeout), m.loading.Init())
 }
 
 func (m Model) View() string {
@@ -45,9 +59,62 @@ func (m Model) View() string {
 
 func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	if !m.connected {
-		var loadscreenCmd tea.Cmd
-		m.loading, loadscreenCmd = m.loading.Update(msg)
-		return m, loadscreenCmd
+		switch msg := msg.(type) {
+		case gateway.ConnectionEstablished:
+			m.id = snowflake.ID(msg)
+			m.connected = true
+			m.timeout = initialTimeout
+			return m, m.timer.Stop()
+
+		case gateway.ConnectionFailed:
+			m.timer = newTimer(m.timeout)
+			m.updateLoadScreenContent()
+			return m, m.timer.Start()
+
+		case timer.TimeoutMsg:
+			m.timeout = min(m.timeout*2, time.Minute)
+			m.loading.SetContent(connectingToServer)
+			return m, gateway.Connect(m.privKey, connectionTimeout)
+
+		case timer.StartStopMsg:
+			var cmd tea.Cmd
+			m.timer, cmd = m.timer.Update(msg)
+			return m, cmd
+
+		case timer.TickMsg:
+			m.updateLoadScreenContent()
+			var cmd tea.Cmd
+			m.timer, cmd = m.timer.Update(msg)
+			return m, cmd
+
+		case spinner.TickMsg:
+			var loadscreenCmd tea.Cmd
+			m.loading, loadscreenCmd = m.loading.Update(msg)
+			return m, loadscreenCmd
+
+		default:
+			return m, nil
+		}
 	}
+
+	switch msg.(type) {
+	case gateway.ConnectionLost:
+		m.connected = false
+		m.timeout = initialTimeout
+		return m, tea.Batch(gateway.Connect(m.privKey, connectionTimeout), m.loading.Init())
+
+	case ui.QuitMsg:
+		gateway.Disconnect()
+	}
+
 	return m, nil
 }
+
+func (m *Model) updateLoadScreenContent() {
+	seconds := m.timer.Timeout.Round(time.Second) / time.Second
+	m.loading.SetContent(fmt.Sprintf(connectionFailed, seconds))
+}
+
+func newTimer(timeout time.Duration) timer.Model {
+	return timer.NewWithInterval(timeout.Truncate(time.Second)+(time.Second/2), timerInterval)
+}

internal/client/ui/ui.go

@@ -15,6 +15,8 @@ import (
 	"github.com/kyren223/eko/pkg/assert"
 )
 
+const DEBUG = true
+
 var Width int
 var Height int
 var BackgroundColor = "#1E1E2E"
@@ -32,6 +34,8 @@ func Transition(model tea.Model) tea.Cmd {
 	}
 }
 
+type QuitMsg struct{}
+
 func AddBorderHeader(header string, headerOffset int, style lipgloss.Style, render string) string {
 	b := style.GetBorderStyle()
 	body := style.UnsetBorderTop().Render(render)

internal/server/server.go

@@ -6,6 +6,7 @@ import (
 	"crypto/rand"
 	"crypto/tls"
 	_ "embed"
+	"encoding/binary"
 	"errors"
 	"fmt"
 	"io"
@@ -159,6 +160,17 @@ func handleConnection(ctx context.Context, conn net.Conn, server *server) {
 	ctx = session.NewContext(ctx, sess)
 	framer := packet.NewFramer()
 
+	// Write ID back, it's useful for the client to know, and signals successful authentication
+	var id [8]byte
+	binary.BigEndian.PutUint64(id[:], uint64(user.ID))
+	_, err = conn.Write(id[:])
+	if err != nil {
+		log.Println(addr, "failed to write user id")
+		conn.Close()
+		log.Println(addr, "disconnected")
+		return
+	}
+
 	defer func() {
 		conn.Close()
 		close(framer.Out)