ygggo/internal/remote/client.go

/*
 *  Copyright (c) 2026 Kaiy Ragur
 *
 *  This Source Code Form is subject to the terms of the Mozilla Public
 *  License, v. 2.0. If a copy of the MPL was not distributed with this
 *  file, You can obtain one at http://mozilla.org/MPL/2.0/.
 */

package remote

import (
	"bufio"
	"crypto/ed25519"
	"encoding/hex"
	"fmt"
	"io"
	"net"
	"time"

	pb "github.com/neilalexander/yggmail/internal/remote/types"
	"github.com/neilalexander/yggmail/internal/transport"
	"github.com/neilalexander/yggmail/internal/utils"
)

type rpcSession struct {
	conn net.Conn
	buf  *bufio.Reader
}

type Client struct {
	publicKey ed25519.PublicKey
	signer    Signer
	transport transport.Transport

	pool chan *rpcSession
}

func NewClient(pubkey ed25519.PublicKey, signer Signer, tr transport.Transport) *Client {
	c := &Client{
		publicKey: pubkey,
		signer:    signer,
		transport: tr,
		pool:      make(chan *rpcSession, 3),
	}
	return c
}

func (c *Client) getConn(target ed25519.PublicKey) (*rpcSession, error) {
	select {
	case session := <-c.pool:
		return session, nil
	default:
		conn, err := c.transport.Dial(hex.EncodeToString(target))
		if err != nil {
			return nil, err
		}
		session := &rpcSession{
			conn: conn,
			buf:  bufio.NewReader(conn),
		}
		return session, nil
	}
}
func (c *Client) returnConn(sess *rpcSession) {
	select {
	case c.pool <- sess:
	default:
		sess.conn.Close()
	}
}

func (c *Client) Call(targetNode ed25519.PublicKey, req *pb.Request, payload io.ReadCloser, size uint32) (*pb.Response, *Payload, error) {
	for range 2 {
		sess, err := c.getConn(targetNode)
		if err != nil {
			return nil, nil, fmt.Errorf("dial failed: %w", err)
		}

		req.AccountPk = utils.EncodeString(c.publicKey)
		req.Timestamp = time.Now().UnixNano()

		if err := SignRequest(req, c.signer); err != nil {
			c.returnConn(sess)
			return nil, nil, err
		}

		resp, respPayload, err := c.doRPC(sess, req, payload, size)
		if err != nil {
			sess.conn.Close()
			return nil, nil, fmt.Errorf("rpc failed: %w", err)
		}

		if respPayload != nil && respPayload.R != nil {
			respPayload.R.Close()
		}

		c.returnConn(sess)
		return resp, respPayload, nil
	}
	return nil, nil, fmt.Errorf("all retries failed")
}

func (c *Client) doRPC(sess *rpcSession, req *pb.Request, payload io.ReadCloser, size uint32) (*pb.Response, *Payload, error) {
	p := &Payload{
		R:    payload,
		Size: size,
	}

	if err := SendFrame(sess.conn, req, p); err != nil {
		return nil, nil, err
	}

	return ReadResponse(sess.buf)
}