raft/client/client.go

package raft_client

import (
	"bufio"
	"errors"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"
)

// Client handles TCP communication with the Raft KV server.
// It is designed to be easily copy-pasted into other projects.
//
// Example Usage:
//
//	client := raft_client.NewClient("127.0.0.1:9011")
//	if err := client.Connect(); err != nil {
//		log.Fatal(err)
//	}
//	defer client.Close()
//
//	// Login
//	if _, err := client.Login("root", "11111"); err != nil {
//		log.Fatal(err)
//	}
//
//	// Set
//	if err := client.Set("mykey", "myvalue"); err != nil {
//		log.Fatal(err)
//	}
//
//	// Get
//	val, err := client.Get("mykey")
//	if err != nil {
//		log.Fatal(err)
//	}
//	fmt.Println("Got:", val)
type Client struct {
	addr    string
	conn    net.Conn
	reader  *bufio.Reader
	writer  *bufio.Writer
	mu      sync.Mutex
	timeout time.Duration
	token   string
}

// NewClient creates a new client instance.
// addr should be in "host:port" format (e.g., "127.0.0.1:9011").
func NewClient(addr string) *Client {
	return &Client{
		addr:    addr,
		timeout: 10 * time.Second,
	}
}

// Connect establishes the TCP connection.
func (c *Client) Connect() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.conn != nil {
		return nil // Already connected
	}

	conn, err := net.DialTimeout("tcp", c.addr, c.timeout)
	if err != nil {
		return err
	}

	c.conn = conn
	c.reader = bufio.NewReader(conn)
	c.writer = bufio.NewWriter(conn)
	return nil
}

// Close closes the connection.
func (c *Client) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.conn != nil {
		err := c.conn.Close()
		c.conn = nil
		c.reader = nil
		c.writer = nil
		return err
	}
	return nil
}

// Login authenticates with the server and stores the session token.
func (c *Client) Login(username, password string) (string, error) {
	// Login command: LOGIN <user> <pass>
	// Note: The server expects an empty line after the command even if no headers are sent.
	resp, err := c.SendRequest(fmt.Sprintf("LOGIN %s %s", username, password), "")
	if err != nil {
		return "", err
	}

	if strings.HasPrefix(resp, "OK ") {
		token := strings.TrimPrefix(resp, "OK ")
		c.mu.Lock()
		c.token = token
		c.mu.Unlock()
		return token, nil
	}
	return "", fmt.Errorf("login failed: %s", resp)
}

// Set stores a key-value pair.
// Uses the body-based SET command for handling values safely.
func (c *Client) Set(key, value string) error {
	// Use Body for value to support special characters/spaces safer
	// Command: SET <key>
	// Body: <value>
	resp, err := c.SendRequest(fmt.Sprintf("SET %s", key), value)
	if err != nil {
		return err
	}
	if resp != "OK" {
		return fmt.Errorf("set failed: %s", resp)
	}
	return nil
}

// SetAsync stores a key-value pair asynchronously (returns OK immediately after proposal).
func (c *Client) SetAsync(key, value string) error {
	resp, err := c.SendRequest(fmt.Sprintf("ASET %s", key), value)
	if err != nil {
		return err
	}
	if resp != "OK" {
		return fmt.Errorf("aset failed: %s", resp)
	}
	return nil
}

// Get retrieves a value by key.
func (c *Client) Get(key string) (string, error) {
	resp, err := c.SendRequest(fmt.Sprintf("GET %s", key), "")
	if err != nil {
		return "", err
	}

	if strings.HasPrefix(resp, "OK ") {
		return strings.TrimPrefix(resp, "OK "), nil
	} else if resp == "ERR not found" {
		return "", errors.New("not found")
	}
	return "", fmt.Errorf("get failed: %s", resp)
}

// Delete removes a key (assuming DELETE command exists or similar, otherwise generic Execute).
// The current server example only shows SET/GET/LOGIN/AUTH/LOGOUT in tcp_server.go.
// If DELETE is supported by the underlying engine but not exposed in TCP switch, this might fail.
// We will omit DELETE specific helper if not sure, but user can use Execute.
// Let's check tcp_server.go... It only has GET, SET, ASET, LOGIN, AUTH, LOGOUT.
// So no DELETE exposed over TCP yet.

// Execute sends a raw command and returns the response line.
// This handles the authentication token injection if logged in.
func (c *Client) Execute(cmd string) (string, error) {
	return c.SendRequest(cmd, "")
}

// SendRequest sends a command with optional body and reads the response.
// It handles the protocol details: Command Line + Headers + Empty Line + Body.
func (c *Client) SendRequest(cmdLine string, body string) (string, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.conn == nil {
		// Auto-reconnect could be implemented here, but for simplicity we fail.
		// Or we can try to connect.
		if err := c.reconnectLocked(); err != nil {
			return "", fmt.Errorf("connection error: %v", err)
		}
	}

	// 1. Send Auth if needed (and not logging in)
	// The current server implementation is stateful (session based).
	// Once authenticated (via LOGIN or AUTH), the session stays authenticated.
	// However, if we just reconnected, we might need to re-auth?
	// The server creates a new Session on connection.
	// So if we have a token, we should probably send AUTH <token> before the command if this is a new connection?
	// But optimizing that away: we assume connection persists.
	// If the user calls Login, we get a token.
	// We can manually call Auth(token) if needed.
	// For simplicity, we assume the user calls Auth() or Login() at start.

	// 2. Send Command
	if _, err := c.writer.WriteString(cmdLine + "\r\n"); err != nil {
		c.conn.Close()
		c.conn = nil
		return "", err
	}

	// 3. Send Headers
	if len(body) > 0 {
		if _, err := c.writer.WriteString(fmt.Sprintf("Content-Length: %d\r\n", len(body))); err != nil {
			c.conn.Close()
			c.conn = nil
			return "", err
		}
	}

	// 4. Send End of Headers (Empty Line)
	if _, err := c.writer.WriteString("\r\n"); err != nil {
		c.conn.Close()
		c.conn = nil
		return "", err
	}

	// 5. Send Body
	if len(body) > 0 {
		if _, err := c.writer.WriteString(body); err != nil {
			c.conn.Close()
			c.conn = nil
			return "", err
		}
	}

	// 6. Flush
	if err := c.writer.Flush(); err != nil {
		c.conn.Close()
		c.conn = nil
		return "", err
	}

	// 7. Read Response
	// The server sends response terminated by \n
	resp, err := c.reader.ReadString('\n')
	if err != nil {
		c.conn.Close()
		c.conn = nil
		return "", err
	}

	return strings.TrimSpace(resp), nil
}

func (c *Client) reconnectLocked() error {
	if c.conn != nil {
		c.conn.Close()
	}
	conn, err := net.DialTimeout("tcp", c.addr, c.timeout)
	if err != nil {
		return err
	}
	c.conn = conn
	c.reader = bufio.NewReader(conn)
	c.writer = bufio.NewWriter(conn)

	// If we have a token, try to restore session
	if c.token != "" {
		// We need to send AUTH command without calling SendRequest to avoid infinite recursion
		// Construct AUTH command
		cmd := fmt.Sprintf("AUTH %s", c.token)
		c.writer.WriteString(cmd + "\r\n\r\n") // Cmd + Empty Line
		c.writer.Flush()
		resp, err := c.reader.ReadString('\n')
		if err == nil && strings.TrimSpace(resp) == "OK" {
			// Auth success
		} else {
			// Auth failed (maybe expired), but connection is open.
			// We don't error out here, just let the actual command fail if it needs auth.
		}
	}

	return nil
}

// Auth restores a session using an existing token.
func (c *Client) Auth(token string) error {
	resp, err := c.SendRequest(fmt.Sprintf("AUTH %s", token), "")
	if err != nil {
		return err
	}
	if resp == "OK" {
		c.mu.Lock()
		c.token = token
		c.mu.Unlock()
		return nil
	}
	return fmt.Errorf("auth failed: %s", resp)
}