refactor: finished refactoring server-side code

2026-07-12 07:19:29 +00:00 · 2024-10-20 20:12:12 +03:00
parent 1997bc8a15
commit 6f83f199f8
7 changed files with 284 additions and 301 deletions
--- a/internal/packet/encoders.go
+++ b/internal/packet/encoders.go
@@ -1,10 +1,9 @@
 package packet

 import (
-	"bytes"
 	"encoding/json"
-	"io"

+	"github.com/kyren223/eko/pkg/assert"
 	"github.com/vmihailenco/msgpack/v5"
 )

@@ -13,7 +12,7 @@ type TypedMessage interface {
 }

 type defaultPacketEncoder struct {
-	io.Reader
+	data       []byte
 	encoding   Encoding
 	packetType PacketType
 }
@@ -26,28 +25,28 @@ func (e defaultPacketEncoder) Type() PacketType {
 	return e.packetType
 }

-func NewJsonEncoder(message TypedMessage) (PacketEncoder, error) {
+func (e defaultPacketEncoder) Payload() []byte {
+	return e.data
+}
+
+func NewJsonEncoder(message TypedMessage) PacketEncoder {
 	data, err := json.Marshal(message)
-	if err != nil {
-		return nil, err
-	}
+	assert.NoError(err, "encoding a message with JSON should never fail")

 	return defaultPacketEncoder{
-		Reader:     bytes.NewReader(data),
+		data:       data,
 		encoding:   EncodingJson,
 		packetType: message.Type(),
-	}, nil
+	}
 }

-func NewMsgPackEncoder(message TypedMessage) (PacketEncoder, error) {
+func NewMsgPackEncoder(message TypedMessage) PacketEncoder {
 	data, err := msgpack.Marshal(message)
-	if err != nil {
-		return nil, err
-	}
+	assert.NoError(err, "encoding a message with msg pack should never fail")

 	return defaultPacketEncoder{
-		Reader:     bytes.NewReader(data),
+		data:       data,
 		encoding:   EncodingMsgPack,
 		packetType: message.Type(),
-	}, nil
+	}
 }
--- a/internal/packet/messages.go
+++ b/internal/packet/messages.go
@@ -4,43 +4,34 @@ import (
 	"github.com/kyren223/eko/internal/data"
 )

-type EkoMessage struct {
-	Message string `msgpack:"message"`
-}
-
-func (m *EkoMessage) Type() PacketType {
-	return TypeEko
-}
-
 type ErrorMessage struct {
 	Error string `msgpack:"error"`
 }

+func NewOkMessage() *ErrorMessage {
+	return &ErrorMessage{}
+}
+
 func (m *ErrorMessage) Type() PacketType {
-	return TypeError
+	return PacketError
 }

-type GetMessagesMessage struct {
-	Since *int64
-	UpTo *int64
+func (m *ErrorMessage) IsOk() bool {
+	return m.Error == ""
 }

-func (m *GetMessagesMessage) Type() PacketType {
-	return TypeGetMessages
-}
-
-type SendMessageMessage struct {
+type SendMessage struct {
 	Content string
 }

-func (m *SendMessageMessage) Type() PacketType {
-	return TypeSendMessage
+func (m *SendMessage) Type() PacketType {
+	return PacketSendMessage
 }

-type MessagesMessage struct {
+type Messages struct {
 	Messages []data.Message
 }

-func (m *MessagesMessage) Type() PacketType {
-	return TypeMessages
+func (m *Messages) Type() PacketType {
+	return PacketMessages
 }
--- a/internal/packet/packet.go
+++ b/internal/packet/packet.go
@@ -60,11 +60,11 @@ func (t PacketType) String() string {
 	case PacketError:
 		return "PacketError"
 	case PacketSendMessage:
-		return "PacketTypeSendMessage"
+		return "PacketSendMessage"
 	case PacketMessages:
-		return "PacketTypeMessages"
+		return "PacketMessages"
 	default:
-		return fmt.Sprintf("PacketTypeInvalid(%v)", byte(t))
+		return fmt.Sprintf("PacketInvalidType(%v)", byte(t))
 	}
 }

@@ -110,15 +110,15 @@ type Packet struct {
 func NewPacket(encoder PacketEncoder) Packet {
 	payload := encoder.Payload()
 	n := len(payload)
-	assert.Assert(0 <= n && n <= PAYLOAD_MAX_SIZE, "size of payload must be valid")
+	assert.Assert(0 <= n && n <= PAYLOAD_MAX_SIZE, "size of payload must be valid", "size", n)

 	data := make([]byte, HEADER_SIZE+n)

 	data[VERSION_OFFSET] = VERSION

 	packetType, encoding := byte(encoder.Type()), byte(encoder.Encoding())
-	assert.Assert(packetType <= 63, "packet type exceeded allowed size type=%v", packetType)
-	assert.Assert(encoding <= 3, "encoding exceeded allowed permutations encoding=%v", encoding)
+	assert.Assert(packetType <= 63, "packet type exceeded allowed size", "type", packetType)
+	assert.Assert(encoding <= 3, "encoding exceeded allowed size", "encoding", encoding)
 	data[TYPE_OFFSET] = packetType | encoding<<6

 	binary.BigEndian.PutUint16(data[LENGTH_OFFSET:], uint16(n))
@@ -170,7 +170,7 @@ func (p Packet) DecodePayload(v TypedMessage) error {
 	case EncodingUnused2:
 		return fmt.Errorf("unsupported encoding: %v", p.Encoding().String())
 	default:
-		assert.Never("encoding from packet should always be valid encoding=%v", p.Encoding())
+		assert.Never("encoding from packet should always be valid", "encoding", p.Encoding())
 		return nil
 	}
 }
--- a/internal/packet/packet_test.go
+++ b/internal/packet/packet_test.go
@@ -1,108 +0,0 @@
-package packet
-
-import (
-	"context"
-	"io"
-	"testing"
-	"time"
-
-	"github.com/vmihailenco/msgpack/v5"
-)
-
-func TestMsgPackEncoding(t *testing.T) {
-	request := EkoMessage{Message: "test"}
-	data, err := msgpack.Marshal(&request)
-	if err != nil {
-		t.Errorf("encoding error: %v", err)
-		return
-	}
-	var response EkoMessage
-	err = msgpack.Unmarshal(data, &response)
-	if err != nil {
-		t.Errorf("decoding error: %v", err)
-		return
-	}
-	if request.Message != response.Message {
-		t.Errorf("%v != %v", request.Message, response.Message)
-		return
-	}
-}
-
-func TestPacketMsgPackEncoding(t *testing.T) {
-	request := EkoMessage{Message: "test"}
-	encoder1, err := NewMsgPackEncoder(&request)
-	encoder2, _ := NewMsgPackEncoder(&request)
-	if err != nil {
-		t.Errorf("encoding error: %v", err)
-		return
-	}
-	encodedBytes := make([]byte, PACKET_MAX_SIZE)
-	n, _ := encoder2.Read(encodedBytes[HEADER_SIZE:])
-
-	packet := NewPacket(encoder1)
-	var response EkoMessage
-	err = packet.DecodePayload(&response)
-	if err != nil {
-		t.Errorf("decoding error: %#v: packet: %v encoder: %v", err, packet.Payload(), encodedBytes[HEADER_SIZE:HEADER_SIZE+n])
-		return
-	}
-	if request.Message != response.Message {
-		t.Errorf("%v != %v", request.Message, response.Message)
-		return
-	}
-}
-
-type TestIoReader struct {
-	data []byte
-	len  int
-}
-
-func (r *TestIoReader) Read(data []byte) (int, error) {
-	if len(r.data) <= r.len {
-		// log.Println("RETURNING EOF:", len(data), r.len)
-		return 0, io.EOF
-	}
-	n := copy(data, r.data[r.len:])
-	r.len += n
-	// log.Println("RETURNING N:", len(r.data), r.len)
-	return n, nil
-}
-
-func (r *TestIoReader) start(msg TypedMessage, t *testing.T) {
-	encoder, err := NewMsgPackEncoder(msg)
-	if err != nil {
-		t.Errorf("encoding error: %v", err)
-		return
-	}
-	r.data = NewPacket(encoder).data
-	// log.Println("LEN:", len(r.data))
-}
-
-func TestPacketFramer(t *testing.T) {
-	reader := &TestIoReader{}
-
-	// Long msg to test multiple
-	msg := EkoMessage{"Testing FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting FramerTesting Framer"}
-	reader.start(&msg, t)
-
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel()
-	out, outErr := RunFramer(ctx, reader)
-
-	var message EkoMessage
-	select {
-	case packet := <-out:
-		if err := packet.DecodePayload(&message); err != nil {
-			t.Errorf("error decoding response: %v", err)
-			return
-		}
-		if msg.Message != message.Message {
-			t.Errorf("%v != %v", msg.Message, message.Message)
-			return
-		}
-
-	case err := <-outErr:
-		t.Errorf("error receiving packet: %v", err)
-		return
-	}
-}
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -2,6 +2,7 @@ package server

 import (
 	"context"
+	"crypto/ed25519"
 	"crypto/rand"
 	"crypto/tls"
 	_ "embed"
@@ -10,6 +11,7 @@ import (
 	"io"
 	"log"
 	"net"
+	"os"
 	"strconv"
 	"strings"
 	"sync"
@@ -48,44 +50,64 @@ func init() {
 	}

 	message := packet.ErrorMessage{Error: packet.PacketUnsupportedEncoding.Error()}
-	encoder, err := packet.NewMsgPackEncoder(&message)
-	assert.NoError(err, "constant packets should not error")
-	unsupportedEncodingErrorPacket = packet.NewPacket(encoder)
+	unsupportedEncodingErrorPacket = packet.NewPacket(packet.NewMsgPackEncoder(&message))

 	message = packet.ErrorMessage{Error: packet.PacketUnsupportedType.Error()}
-	encoder, err = packet.NewMsgPackEncoder(&message)
-	assert.NoError(err, "constant packets should not error")
-	unsupportedTypeErrorPacket = packet.NewPacket(encoder)
+	unsupportedTypeErrorPacket = packet.NewPacket(packet.NewMsgPackEncoder(&message))
 }

 type server struct {
-	node *snowflake.Node
-	port uint16
+	Node     *snowflake.Node
+	Port     uint16
+	sessions map[snowflake.ID]*Session
+	sessMu   sync.RWMutex
 }

 func NewServer(port uint16) server {
-	assert.Assert(nodeId <= snowflake.NodeMax, "maximum amount of servers reached %v", snowflake.NodeMax)
+	assert.Assert(nodeId <= snowflake.NodeMax, "maximum amount of servers reached")
 	node := snowflake.NewNode(nodeId)
 	nodeId++

 	return server{
-		node: node,
-		port: port,
+		Node:     node,
+		Port:     port,
+		sessions: map[snowflake.ID]*Session{},
 	}
 }

+func (s *server) AddSession(session *Session) {
+	s.sessMu.Lock()
+	defer s.sessMu.Unlock()
+	s.sessions[session.ID] = session
+}
+
+func (s *server) RemoveSession(id snowflake.ID) {
+	s.sessMu.Lock()
+	defer s.sessMu.Unlock()
+	delete(s.sessions, id)
+}
+
+func (s *server) Session(id snowflake.ID) (*Session, bool) {
+	s.sessMu.RLock()
+	defer s.sessMu.RUnlock()
+	session, ok := s.sessions[id]
+	return session, ok
+}
+
 func (s *server) ListenAndServe(ctx context.Context) {
-	listener, err := tls.Listen("tcp4", ":"+strconv.Itoa(int(s.port)), tlsConfig)
+	listener, err := tls.Listen("tcp4", ":"+strconv.Itoa(int(s.Port)), tlsConfig)
 	if err != nil {
 		log.Fatalf("error starting server: %s", err)
 	}
+
+	assert.AddFlush(listener)
 	defer listener.Close()
 	go func() {
 		<-ctx.Done()
 		listener.Close()
 	}()

-	log.Printf("started listening on port %v...\n", s.port)
+	log.Printf("started listening on port %v...\n", s.Port)
 	var wg sync.WaitGroup
 	for {
 		conn, err := listener.Accept()
@@ -97,46 +119,61 @@ func (s *server) ListenAndServe(ctx context.Context) {
 		}
 		wg.Add(1)
 		go func() {
-			handleConnection(ctx, conn)
+			handleConnection(ctx, conn, s)
 			wg.Done()
 		}()
 	}
-	log.Printf("stopped listening on port %v\n", s.port)
+	log.Printf("stopped listening on port %v\n", s.Port)

 	log.Println("waiting for all active connections to close...")
 	wg.Wait()
 	log.Println("server shutdown complete")
 }

-func handleConnection(ctx context.Context, conn net.Conn) {
+func handleConnection(ctx context.Context, conn net.Conn, server *server) {
 	addr, ok := conn.RemoteAddr().(*net.TCPAddr)
 	assert.Assert(ok, "getting tcp address should never fail as we are using tcp connections")

-	writeQueue := make(chan packet.Packet, 10)
-	session := newSession(addr, writeQueue)
-	nonce := session.Challenge()
+	log.Println(addr, "accepted")

+	nonce := [32]byte{}
+	_, err := rand.Read(nonce[:])
+	assert.NoError(err, "random should always produce a value")
+	pubKey, err := handleAuth(conn, nonce[:])
+	if err != nil {
+		log.Println(addr, err)
+		conn.Close()
+		log.Println(addr, "disconnected")
+		return
+	}
+
+	// TODO: replace this with DB query for id
+	id := server.Node.Generate()
+	session := newSession(server, addr, id, pubKey)
+	server.AddSession(session)
 	ctx = newContext(ctx, session)
 	framer := packet.NewFramer(ctx)

-	log.Println(addr, "accepted")
-
 	defer func() {
 		conn.Close()
+		close(framer.Out)
+		server.RemoveSession(session.ID)
+		close(session.WriteQueue)
 		log.Println(addr, "disconnected")
 	}()

 	go func() {
-		var mu sync.Mutex
 		for {
-			packet, ok := <-writeQueue
+			packet, ok := <-session.WriteQueue
 			if !ok {
 				break
 			}
-			mu.Lock()
-			packet.Into(conn)
-			mu.Unlock()
+			if _, err := packet.Into(conn); err != nil {
+				log.Println(addr, err)
+				break
+			}
 		}
+		session.WriteQueue = nil
 	}()

 	go func() {
@@ -146,91 +183,97 @@ func handleConnection(ctx context.Context, conn net.Conn) {
 				break
 			}
 			response := processPacket(ctx, request)
-			writeQueue <- response
+			if session.WriteQueue == nil {
+				break
+			}
+			session.WriteQueue <- response
 		}
 	}()

 	buffer := make([]byte, 512)
 	for {
+		conn.SetReadDeadline(time.Now().Add(time.Second))
 		n, err := conn.Read(buffer)
-		if err != nil {
+		deadlineExceeded := errors.Is(err, os.ErrDeadlineExceeded)
+		if err != nil && !deadlineExceeded {
 			if !errors.Is(err, io.EOF) {
 				log.Println(addr, "read error:", err)
 			}
 			break
 		}

+		if ctx.Err() != nil {
+			log.Println(addr, ctx.Err())
+			break
+		}
+
 		err = framer.Push(ctx, buffer[:n])
 		if err != nil {
-			// Wrap err and send to client then break
+			if ctx.Err() != nil {
+				log.Println(addr, ctx.Err())
+			} else {
+				// TODO: Wrap err and send to client then break
+			}
+			break
 		}
 	}
 }

-func _handleConnection(ctx context.Context, conn net.Conn) {
-	addr, ok := conn.RemoteAddr().(*net.TCPAddr)
-	assert.Assert(ok, "getting tcp address should be valid")
+func handleAuth(conn net.Conn, nonce []byte) (ed25519.PublicKey, error) {
+	conn.SetDeadline(time.Now().Add(time.Second * 5))

-	log.Println(addr, "accepted")
-	defer log.Println(addr, "disconnected")
-	defer conn.Close()
+	challengePacket := make([]byte, len(nonce)+1)
+	challengePacket[0] = packet.VERSION
+	copy(challengePacket[1:], nonce)

-	// ctx = newContext(ctx, addr)
-	// // TODO: consider adding timeout/deadline to ctx?
-
-	out, outErr := packet.RunFramer(ctx, conn)
-outer:
-	for {
-		select {
-		case packet, ok := <-out:
-			if !ok {
-				break outer
-			}
-			log.Printf("client %v: request packet: %v\n", conn.RemoteAddr().String(), packet)
-			responsePacket, err := handlePacket(packet)
-			log.Printf("client %v: response packet: %v\n", conn.RemoteAddr().String(), responsePacket)
-			if err != nil {
-				log.Printf("client %v: error processing request: %v\n", conn.RemoteAddr().String(), err)
-				break outer
-			}
-			_, err = responsePacket.Into(conn)
-			if err != nil {
-				log.Printf("client %v: error writing packet: %v\n", conn.RemoteAddr().String(), err)
-				break outer
-			}
-
-		case err := <-outErr:
-			if err == packet.PacketUnsupportedEncoding {
-				_, err := unsupportedEncodingErrorPacket.Into(conn)
-				log.Printf("client %v: error writing unsupported encoding packet: %v\n", conn.RemoteAddr().String(), err)
-			} else if err == packet.PacketUnsupportedType {
-				_, err := unsupportedTypeErrorPacket.Into(conn)
-				log.Printf("client %v: error writing unsupported type packet: %v\n", conn.RemoteAddr().String(), err)
-			} else if err != nil {
-				log.Printf("client %v: internal error: %v\n", conn.RemoteAddr().String(), err)
-			}
-			break outer
-
-		case <-ctx.Done():
-			log.Printf("client %v: %v\n", conn.RemoteAddr().String(), ctx.Err())
-			break outer
-		}
+	_, 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")
+	}
+
+	conn.SetDeadline(time.Time{})
+	return pubKey, nil
 }

 type Session struct {
+	Server     *server
 	Addr       *net.TCPAddr
-	WriteQueue <-chan packet.Packet
+	WriteQueue chan packet.Packet
+	ID         snowflake.ID
+	PubKey     ed25519.PublicKey

 	mu         sync.Mutex
-	challenge   []byte
+	challenge  []byte
 	issuedTime time.Time
 }

-func newSession(addr *net.TCPAddr, writeQueue <-chan packet.Packet) *Session {
+func newSession(server *server, addr *net.TCPAddr, id snowflake.ID, pubKey ed25519.PublicKey) *Session {
 	session := &Session{
+		Server:     server,
 		Addr:       addr,
-		WriteQueue: writeQueue,
+		PubKey:     pubKey,
+		WriteQueue: make(chan packet.Packet, 10),
 		challenge:  make([]byte, 32), // Recommended nonce size
 	}
 	session.Challenge() // Make sure an initial nonce is generated
@@ -263,58 +306,43 @@ func FromContext(ctx context.Context) (*Session, bool) {

 // TODO: Move everything below this to somewhere else

-func handlePacket(pkt packet.Packet) (packet.Packet, error) {
+func processPacket(ctx context.Context, pkt packet.Packet) packet.Packet {
+	session, ok := FromContext(ctx)
+	assert.Assert(ok, "context in process packet should always have a session")
+
 	var response packet.TypedMessage
 	switch pkt.Type() {
-	case packet.PacketEko:
-		var request packet.EkoMessage
-		if err := pkt.DecodePayload(&request); err != nil {
-			return packet.Packet{}, fmt.Errorf("decode error: %v", err)
-		}
-
-		response = &packet.EkoMessage{Message: "Eko \"" + request.Message + "\""}
 	case packet.PacketSendMessage:
-		var request packet.SendMessageMessage
+		var request packet.SendMessage
 		if err := pkt.DecodePayload(&request); err != nil {
-			return packet.Packet{}, fmt.Errorf("decode error: %v", err)
+			log.Println("decode error:", err)
+			response = &packet.ErrorMessage{Error: "malformed payload"}
+			break
 		}

 		content := strings.TrimSpace(request.Content)
 		if content == "" {
-			response = &packet.ErrorMessage{Error: "content must not be blank"}
+			response = &packet.ErrorMessage{Error: "message content must not be blank"}
 			break
 		}

+		node := session.Server.Node
 		message := data.Message{
 			Id:          node.Generate(),
-			SenderId:    node.Generate(),
-			FrequencyId: node.Generate(),
-			NetworkId:   node.Generate(),
+			SenderId:    session.ID,
+			FrequencyId: node.Generate(), // TODO: replace with actual ID
+			NetworkId:   node.Generate(), // TODO: replace with actual ID
 			Contents:    content,
 		}
 		messages = append(messages, message)

-		response = &packet.EkoMessage{Message: "Eko OK"}
-	case packet.PacketGetMessages:
-		var request packet.GetMessagesMessage
-		if err := pkt.DecodePayload(&request); err != nil {
-			return packet.Packet{}, fmt.Errorf("decode error: %v", err)
-		}
-
-		response = &packet.MessagesMessage{Messages: messages}
+		response = packet.NewOkMessage()
 	default:
-		return packet.Packet{}, errors.New("TODO: not implemented yet")
+		response = &packet.ErrorMessage{Error: "use of unsupported packet type"}
 	}

-	assert.NotNil(response, "response must always be set")
-	encoder, err := packet.NewMsgPackEncoder(response)
-	if err != nil {
-		return packet.Packet{}, fmt.Errorf("encode error: %v", err)
-	}
-	return packet.NewPacket(encoder), nil
+	assert.NotNil(response, "response must always be assigned to")
+	return packet.NewPacket(packet.NewMsgPackEncoder(response))
 }

-var (
-	node     = snowflake.NewNode(1)
-	messages []data.Message
-)
+var messages []data.Message
--- a/pkg/assert/assert.go
+++ b/pkg/assert/assert.go
@@ -1,25 +1,94 @@
 package assert

-import "log"
+import (
+	"fmt"
+	"io"
+	"os"
+	"reflect"
+	"runtime/debug"
+	"sync"
+)

-func Assert(assertion bool, message string, a ...any) {
+var (
+	writer io.Writer = os.Stderr
+
+	flushes = []io.Closer{}
+	flushMu sync.Mutex
+
+	assertData = map[string]any{}
+	mapMu      sync.Mutex
+)
+
+func AddData(key string, value any) {
+	mapMu.Lock()
+	assertData[key] = value
+	mapMu.Unlock()
+}
+
+func RemoveData(key string) {
+	mapMu.Lock()
+	delete(assertData, key)
+	mapMu.Unlock()
+}
+
+func AddFlush(flusher io.Closer) {
+	flushMu.Lock()
+	flushes = append(flushes, flusher)
+	flushMu.Unlock()
+}
+
+func SetWriter(w io.Writer) {
+	writer = w
+}
+
+func runAssert(message string, args ...any) {
+	flushMu.Lock()
+	for len(flushes) != 0 {
+		flusher := flushes[len(flushes)-1]
+		flusher.Close()
+		flushes = flushes[:len(flushes)-1]
+	}
+	flushMu.Unlock()
+
+	values := []any{
+		"msg", message,
+	}
+	values = append(values, args...)
+	mapMu.Lock()
+	for k, v := range assertData {
+		values = append(values, k, v)
+	}
+	mapMu.Unlock()
+
+	fmt.Fprintf(writer, "ARGS: %+v\n", args)
+	fmt.Fprintf(writer, "ASSERT\n")
+	for i := 0; i < len(values); i += 2 {
+		fmt.Fprintf(writer, "   %s=%v\n", values[i], values[i+1])
+	}
+	fmt.Fprintln(writer, string(debug.Stack()))
+
+	os.Exit(1)
+}
+
+func Assert(assertion bool, message string, args ...any) {
 	if !assertion {
-		log.Fatalf(message+"\n", a...)
+		runAssert(message, args...)
 	}
 }

-func NoError(err error, message string, a ...any) {
+func NoError(err error, message string, args ...any) {
 	if err != nil {
-		log.Fatalf(message+": "+err.Error()+"\n", a...)
+		args = append(args, "error", err)
+		runAssert(message, args...)
 	}
 }

-func Never(message string, a ...any) {
-	log.Fatalf(message+"\n", a...)
+func Never(message string, args ...any) {
+	runAssert(message, args...)
 }

-func NotNil(value any, message string, a ...any) {
-	if value == nil {
-		log.Fatalf(message+"\n", a...)
+func NotNil(value any, message string, args ...any) {
+	if value == nil || reflect.ValueOf(value).Kind() == reflect.Ptr && reflect.ValueOf(value).IsNil() {
+		runAssert(message, args...)
 	}
 }
--- a/pkg/snowflake/snowflake.go
+++ b/pkg/snowflake/snowflake.go
@@ -12,16 +12,34 @@ import (
 const (
 	// Epoch is set to the twitter snowflake epoch of Nov 04 2010 01:42:54 UTC in milliseconds
 	// TODO: change this to eko epoch when eko is production ready
-	Epoch int64 = 1288834974657
-	nodeBits  = 10
-	stepBits  = 12
-	NodeMax   = 1<<nodeBits - 1
-	nodeMask  = NodeMax << stepBits
-	stepMask  = 1<<stepBits - 1
-	timeShift = nodeBits + stepBits
-	nodeShift = stepBits
+	Epoch     int64 = 1288834974657
+	nodeBits        = 10
+	stepBits        = 12
+	NodeMax         = 1<<nodeBits - 1
+	nodeMask        = NodeMax << stepBits
+	stepMask        = 1<<stepBits - 1
+	timeShift       = nodeBits + stepBits
+	nodeShift       = stepBits
 )

+type ID int64
+
+func (id ID) String() string {
+	return strconv.FormatInt(int64(id), 10)
+}
+
+func (id ID) Time() (msSinceEpoch int64) {
+	return (int64(id) >> timeShift) + Epoch
+}
+
+func (id ID) Node() int64 {
+	return int64(id) & nodeMask >> nodeShift
+}
+
+func (id ID) Step() int64 {
+	return int64(id) & stepMask
+}
+
 type Node struct {
 	mu    sync.Mutex
 	epoch time.Time
@@ -30,11 +48,13 @@ type Node struct {
 	step  int64
 }

-type ID int64
+func (n *Node) String() string {
+	return fmt.Sprintf("Node%v(epoch: %v, step: %v, time: %v)", n.node, n.epoch, n.step, n.time)
+}

 func NewNode(node int64) *Node {
-	assert.Assert(nodeBits+stepBits <= 22, "node and step bits must add up to 22 or less")
-	assert.Assert(0 <= node && node <= NodeMax, "node and step bits must add up to 22 or less")
+	assert.Assert(nodeBits+stepBits == 22, "node and step bits must add up to 22")
+	assert.Assert(0 <= node && node <= NodeMax, "node must be within 0 and NodeMax", "node", node)

 	// Credit to https://github.com/bwmarrin/snowflake
 	currentTime := time.Now()
@@ -67,22 +87,6 @@ func (n *Node) Generate() ID {
 	return ID((now << timeShift) | (n.node << nodeShift) | (n.step))
 }

-func (id ID) String() string {
-	return strconv.FormatInt(int64(id), 10)
-}
-
-func (id ID) Time() (msSinceEpoch int64) {
-	return (int64(id) >> timeShift) + Epoch
-}
-
-func (id ID) Node() int64 {
-	return int64(id) & nodeMask >> nodeShift
-}
-
-func (id ID) Step() int64 {
-	return int64(id) & stepMask
-}
-
 // Json marshling to avoid inprecision of json number (float64)

 func (id ID) MarshalJSON() ([]byte, error) {