Refactored gateway to integrate between with the UI

6 files changed, 118 insertions, 79 deletions

diff --git a/internal/client/client.go b/internal/client/client.go
index 938a265..d658427 100644
--- a/internal/client/client.go
+++ b/internal/client/client.go
@@ -38,11 +38,8 @@ func Run() {
 
 	program := tea.NewProgram(initialModel(), tea.WithAltScreen())
 	assert.AddFlush(BubbleTeaCloser{program})
+	ui.Program = program
 
-	// _, privKey, err := ed25519.GenerateKey(nil)
-	// assert.NoError(err, "private key gen should not error")
-
-	// gateway.Connect(context.Background(), program, privKey)
 	if _, err := program.Run(); err != nil {
 		log.Println(err)
 	}
diff --git a/internal/client/gateway/gateway.go b/internal/client/gateway/gateway.go
index 7e9ab82..441395a 100644
--- a/internal/client/gateway/gateway.go
+++ b/internal/client/gateway/gateway.go
@@ -6,6 +6,8 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	_ "embed"
+	"errors"
+	"io"
 	"log"
 	"net"
 	"sync"
@@ -13,6 +15,7 @@ import (
 
 	tea "github.com/charmbracelet/bubbletea"
 
+	"github.com/kyren223/eko/internal/client/ui"
 	"github.com/kyren223/eko/internal/packet"
 	"github.com/kyren223/eko/pkg/assert"
 )
@@ -31,6 +34,13 @@ var (
 	writeMu sync.Mutex
 )
 
+type (
+	ConnectionEstablished struct{}
+	ConnectionFailed      error
+	ConnectionLost        error
+	ConnectionClosed      struct{}
+)
+
 func init() {
 	certPool := x509.NewCertPool()
 	if !certPool.AppendCertsFromPEM(certPEM) {
@@ -43,7 +53,17 @@ func init() {
 	}
 }
 
-func Connect(ctx context.Context, program *tea.Program, privKey ed25519.PrivateKey) error {
+func Connect(ctx context.Context, privKey ed25519.PrivateKey) tea.Cmd {
+	return func() tea.Msg {
+		err := connect(ctx, privKey)
+		if err != nil {
+			return ConnectionFailed(err)
+		}
+		return ConnectionEstablished{}
+	}
+}
+
+func connect(ctx context.Context, privKey ed25519.PrivateKey) error {
 	assert.Assert(conn == nil, "cannot connect, connection is active")
 
 	connChan := make(chan net.Conn, 1)
@@ -74,25 +94,12 @@ func Connect(ctx context.Context, program *tea.Program, privKey ed25519.PrivateK
 		return ctx.Err()
 	}
 
-	go readUntilDisconnected()
-	go handlePacketStream(program)
+	go readForever()
+	go handlePacketStream()
 
 	return nil
 }
 
-func Disconnect() {
-	assert.Assert(conn != nil, "cannot disconnect, connection is inactive")
-	close(framer.Out)
-	conn.Close()
-	conn = nil
-	responsesMu.Lock()
-	for _, responseChan := range asyncResponses {
-		close(responseChan)
-	}
-	asyncResponses = nil
-	responsesMu.Unlock()
-}
-
 func handleAuth(ctx context.Context, privKey ed25519.PrivateKey) error {
 	const nonceSize = 32
 	challengeRequest := make([]byte, 1+nonceSize)
@@ -131,41 +138,31 @@ func handleAuth(ctx context.Context, privKey ed25519.PrivateKey) error {
 	return nil
 }
 
-func readUntilDisconnected() {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-	go func() {
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			default:
-				if conn == nil {
-					cancel()
-					return
-				}
-			}
-		}
-	}()
-
+func readForever() {
 	buffer := make([]byte, 512)
 	for conn != nil {
 		n, err := conn.Read(buffer)
 		if err != nil {
-			log.Println("server connectivity error: ", err)
-			break
+			if errors.Is(err, io.EOF) {
+				err = nil
+			}
+			onDisconnect(err)
+			return
 		}
 
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 		err = framer.Push(ctx, buffer[:n])
 		if ctx.Err() != nil {
-			log.Println("server connectivity error: ", ctx.Err())
-			break
+			cancel()
+			onDisconnect(errors.New("framer blocked for more than a second, closing connection"))
+			return
 		}
+		cancel()
 		assert.NoError(err, "packets from server should always be correctly formatted")
 	}
 }
 
-func handlePacketStream(program *tea.Program) {
+func handlePacketStream() {
 	for {
 		pkt, ok := <-framer.Out
 		if !ok {
@@ -177,7 +174,7 @@ func handlePacketStream(program *tea.Program) {
 
 		if pkt.Type().IsPush() {
 			log.Println("received streamed packet:", payload)
-			program.Send(payload)
+			ui.Program.Send(payload)
 			continue
 		}
 
@@ -194,8 +191,32 @@ func handlePacketStream(program *tea.Program) {
 	}
 }
 
+func Disconnect() {
+	assert.Assert(conn != nil, "cannot disconnect, connection is inactive")
+	conn.Close()
+}
+
+func onDisconnect(err error) {
+	conn.Close()
+	conn = nil
+	close(framer.Out)
+	responsesMu.Lock()
+	for _, responseChan := range asyncResponses {
+		close(responseChan)
+	}
+	asyncResponses = nil
+	responsesMu.Unlock()
+	if err != nil {
+		log.Println("connection lost:", err)
+		ui.Program.Send(ConnectionLost(err))
+	} else {
+		log.Println("connection closed")
+		ui.Program.Send(ConnectionClosed{})
+	}
+}
+
 func Send(request packet.Payload) <-chan packet.Payload {
-	responseChan := make(chan packet.Payload)
+	responseChan := make(chan packet.Payload, 1)
 	go func() {
 		pkt := packet.NewPacket(packet.NewMsgPackEncoder(request))
 
@@ -205,7 +226,6 @@ func Send(request packet.Payload) <-chan packet.Payload {
 
 		if conn == nil {
 			log.Println("request send error:", "connection is closed")
-			close(responseChan)
 			return
 		}
 
@@ -214,7 +234,6 @@ func Send(request packet.Payload) <-chan packet.Payload {
 		writeMu.Unlock()
 		if err != nil {
 			log.Println("request send error:", err)
-			close(responseChan)
 			return
 		}
 	}()
diff --git a/internal/client/ui/auth/auth.go b/internal/client/ui/auth/auth.go
index aa988c5..2e81a16 100644
--- a/internal/client/ui/auth/auth.go
+++ b/internal/client/ui/auth/auth.go
@@ -20,7 +20,7 @@ import (
 	"github.com/kyren223/eko/internal/client/ui"
 	authfield "github.com/kyren223/eko/internal/client/ui/auth/field"
 	"github.com/kyren223/eko/internal/client/ui/choicepopup"
-	"github.com/kyren223/eko/internal/client/ui/loadscreen"
+	"github.com/kyren223/eko/internal/client/ui/core"
 	"github.com/kyren223/eko/pkg/assert"
 )
 
@@ -474,7 +474,8 @@ func (m *Model) Signup() tea.Cmd {
 		return nil
 	}
 
-	return authenticate(privKey)
+
+	return ui.Transition(core.New(privKey, username))
 }
 
 func (m *Model) signin() tea.Cmd {
@@ -538,7 +539,8 @@ func (m *Model) signin() tea.Cmd {
 			config.PrivateKeyPath = privateKeyFilepath
 		})
 	}
-	return authenticate(*privKey)
+
+	return ui.Transition(core.New(*privKey, ""))
 }
 
 func (m Model) ButtonIndex() int {
@@ -579,10 +581,3 @@ func expandPath(path string) string {
 	return path
 }
 
-func authenticate(privKey ed25519.PrivateKey) tea.Cmd {
-	text := lipgloss.NewStyle().Inline(true).
-		Render("Update Failed - retrying in 3 sec...")
-	return ui.Transition(loadscreen.New(text, func(msg tea.Msg) tea.Cmd {
-		return nil
-	}))
-}
diff --git a/internal/client/ui/core/core.go b/internal/client/ui/core/core.go
new file mode 100644
index 0000000..1dbaef5
--- /dev/null
+++ b/internal/client/ui/core/core.go
@@ -0,0 +1,36 @@
+package core
+
+import (
+	"crypto/ed25519"
+	"fmt"
+
+	tea "github.com/charmbracelet/bubbletea"
+)
+
+type Model struct{}
+
+func New(privKey ed25519.PrivateKey, name string) Model {
+
+	// Connectiong to server...
+	// Connection failed - retrying in 3 sec...
+	// Connection failed - retrying in 3 sec...
+	return Model{}
+}
+
+func (m Model) Init() tea.Cmd {
+	return nil
+}
+
+func (m Model) View() string {
+	return fmt.Sprintf(
+		"%s",
+		"",
+	)
+}
+
+func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	switch msg := msg.(type) {
+	default:
+		return m, nil
+	}
+}
diff --git a/internal/client/ui/loadscreen/loadscreen.go b/internal/client/ui/loadscreen/loadscreen.go
index 869237f..6f4610a 100644
--- a/internal/client/ui/loadscreen/loadscreen.go
+++ b/internal/client/ui/loadscreen/loadscreen.go
@@ -117,18 +117,14 @@ func circleTrail(width int, height int, offset int, clockwise bool, bg string, t
 type Updater func(msg tea.Msg) tea.Cmd
 
 type Model struct {
-	sp      spinner.Model
 	content string
-	updater Updater
+	sp      spinner.Model
 }
 
-func New(content string, updater Updater) Model {
-	width := lipgloss.Width(content)
-	height := lipgloss.Height(content)
-	content = lipgloss.NewStyle().Width(width).Height(height).Render(content)
+func New(content string) Model {
 	return Model{
 		sp:      spinner.New(spinner.WithSpinner(loading)),
-		content: style.Render(content),
+		content: content,
 	}
 }
 
@@ -137,28 +133,23 @@ func (m Model) Init() tea.Cmd {
 }
 
 func (m Model) View() string {
+	width := lipgloss.Width(m.content)
+	height := lipgloss.Height(m.content)
+	content := lipgloss.NewStyle().Width(width).Height(height).Render(m.content)
+	content = style.Render(content)
 	return lipgloss.Place(
 		ui.Width, ui.Height,
 		lipgloss.Center, lipgloss.Center,
-		lipgloss.JoinVertical(lipgloss.Center, m.sp.View(), m.content),
+		lipgloss.JoinVertical(lipgloss.Center, m.sp.View(), content),
 	)
 }
 
 func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
-	switch msg := msg.(type) {
-
-	case tea.KeyMsg:
-		key := msg.Type
-		switch key {
-		case tea.KeyCtrlC:
-			return m, tea.Quit
-		}
-
-	case spinner.TickMsg:
-		var cmd tea.Cmd
-		m.sp, cmd = m.sp.Update(msg)
-		return m, cmd
-	}
+	var cmd tea.Cmd
+	m.sp, cmd = m.sp.Update(msg)
+	return m, cmd
+}
 
-	return m, m.updater(msg)
+func (m *Model) SetContent(content string) {
+	m.content = content
 }
diff --git a/internal/client/ui/ui.go b/internal/client/ui/ui.go
index a1db45f..8cda302 100644
--- a/internal/client/ui/ui.go
+++ b/internal/client/ui/ui.go
@@ -20,6 +20,7 @@ var Height int
 var BackgroundColor = "#1E1E2E"
 var Background = lipgloss.Color(BackgroundColor)
 var Style = lipgloss.NewStyle().Background(Background)
+var Program *tea.Program
 
 type ModelTransition struct {
 	Model tea.Model