From 00dd876f252fb9bcdbfd0c7ffc0eaebcdde661dd Mon Sep 17 00:00:00 2001 From: Kyren223 Date: Fri, 22 Nov 2024 16:46:19 +0200 Subject: [PATCH] Implemented proper connection and retry mechanisms --- cmd/client/main.go | 2 +- cmd/server/main.go | 2 +- internal/client/client.go | 33 +++++++++++-- internal/client/gateway/gateway.go | 53 ++++++++++++-------- internal/client/ui/core/core.go | 79 +++++++++++++++++++++++++++--- internal/client/ui/ui.go | 4 ++ internal/server/server.go | 12 +++++ 7 files changed, 153 insertions(+), 32 deletions(-) diff --git a/cmd/client/main.go b/cmd/client/main.go index 61ea570..5347554 100644 --- a/cmd/client/main.go +++ b/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) } diff --git a/cmd/server/main.go b/cmd/server/main.go index 2f83b48..0359ff8 100644 --- a/cmd/server/main.go +++ b/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) } diff --git a/internal/client/client.go b/internal/client/client.go index 57ee5ab..ec760dd 100644 --- a/internal/client/client.go +++ b/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 } diff --git a/internal/client/gateway/gateway.go b/internal/client/gateway/gateway.go index 22d851f..9151d74 100644 --- a/internal/client/gateway/gateway.go +++ b/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)) } } diff --git a/internal/client/ui/core/core.go b/internal/client/ui/core/core.go index 12b8b42..790c52e 100644 --- a/internal/client/ui/core/core.go +++ b/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) +} diff --git a/internal/client/ui/ui.go b/internal/client/ui/ui.go index 8cda302..762d0b6 100644 --- a/internal/client/ui/ui.go +++ b/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) diff --git a/internal/server/server.go b/internal/server/server.go index 07cd537..dce8938 100644 --- a/internal/server/server.go +++ b/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)