ygggo/internal/remote/server.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"
	"context"
	"crypto/ed25519"
	"fmt"
	"io"
	"log"
	"net"
	"time"

	"github.com/neilalexander/yggmail/internal/account"
	"github.com/neilalexander/yggmail/internal/notify"
	pb "github.com/neilalexander/yggmail/internal/remote/types"
	"github.com/neilalexander/yggmail/internal/storage"
	"github.com/neilalexander/yggmail/internal/utils"
)

type Server struct {
	log           *log.Logger
	notify        Notify
	queue         Queues
	accMgr        *account.Manager
	storage       storage.Storage
	nodePublickey ed25519.PublicKey
	validator     *ReqValidator
}

type Notify interface {
	Subscribe(pbk string) chan notify.Event
	Unsubscribe(pbk string, ch chan notify.Event)
	NotifyNew(pk string, ev notify.Event) error
}

type Queues interface {
	QueueForServer(serverHex string) error
}

func NewRemoteServer(nodePublickey ed25519.PublicKey, s storage.Storage, m *account.Manager, l *log.Logger, notify Notify, queue Queues) *Server {
	return &Server{
		storage:       s,
		accMgr:        m,
		log:           l,
		notify:        notify,
		queue:         queue,
		nodePublickey: nodePublickey,
		validator:     NewValidator(time.Minute * 3),
	}
}

func (srv *Server) Serve(l net.Listener) error {
	srv.log.Println("Unified Remote Storage RPC server started (Streaming Mode)")
	for {
		conn, err := l.Accept()
		if err != nil {
			continue
		}
		// Serving in new gorutine
		go srv.handleSession(conn)
	}
}

func (srv *Server) handleSession(conn net.Conn) {
	defer conn.Close()
	br := bufio.NewReader(conn)
	for {
		conn.SetDeadline(time.Now().Add(30 * time.Second))

		req, reqPayload, err := ReadRequest(br)
		if err != nil {
			if err != io.EOF {
				srv.log.Printf("RPC: Session error: %v", err)
			}
			return
		}

		func() {
			conn.SetDeadline(time.Now().Add(2 * time.Minute))
			ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(2*time.Minute))
			defer cancel()

			resp, respPayload := srv.handleRequest(ctx, req, reqPayload)

			err = SendFrame(conn, resp, respPayload)

			if reqPayload != nil {
				io.Copy(io.Discard, reqPayload.R)
				reqPayload.Close()
			}
			if respPayload != nil {
				respPayload.Close()
			}

			if err != nil {
				srv.log.Printf("RPC: Send error: %v", err)
				return
			}
		}()
	}
}
func RespErr(err error) *pb.Response {
	return &pb.Response{Success: false, Error: err.Error()}
}

func (srv *Server) handleRequest(ctx context.Context, req *pb.Request, reqPayload *Payload) (resp *pb.Response, respPayload *Payload) {
	// srv.log.Printf("RPC: Request received from %s", req.AccountPk)
	sPk, err := utils.DecodeString(req.AccountPk)
	if err != nil {
		return RespErr(fmt.Errorf("Error decode key")), respPayload
	}

	// if ok, err := VerifyRequest(req, sPk); err != nil || !ok {
	// 	return nil, fmt.Errorf("401 Unauthorized: Signature mismatch")
	// }
	// srv.log.Printf("RPC: Request Verifed from %s", req.AccountPk)
	ok, err := srv.validator.Validate(req, sPk)
	if err != nil {
		return RespErr(fmt.Errorf("401 Unauthorized: Signature mismatch")), respPayload
	}
	if !ok {
		return &pb.Response{Success: true, Error: "duplicate request ignored"}, respPayload
	}

	// Check if this AccountPubKey belongs to a local account we manage
	acc, _ := srv.accMgr.GetAccount(req.AccountPk)
	isLocalAccount := (acc != nil)

	_, isPushCall := req.Call.(*pb.Request_MailPush)
	if !isLocalAccount && !isPushCall {
		return RespErr(fmt.Errorf("403 Forbidden: Remote management not allowed")), respPayload
	}

	var rcname string
	switch call := req.Call.(type) {
	// --- Mailbox management (Local Account only via Policy) ---
	case *pb.Request_MboxList:
		resp, err = srv.handleMailboxList(ctx, req.AccountPk, call.MboxList)
		rcname = "MboxList"
	case *pb.Request_MboxCreate:
		resp, err = srv.handleMailboxCreate(ctx, req.AccountPk, call.MboxCreate)
		rcname = "MboxCreate"
	case *pb.Request_MboxDelete:
		resp, err = srv.handleMailboxDelete(ctx, req.AccountPk, call.MboxDelete)
		rcname = "MboxDelete"
	case *pb.Request_MboxRename:
		resp, err = srv.handleMailboxRename(ctx, req.AccountPk, call.MboxRename)
		rcname = "MboxRename"
	case *pb.Request_MboxSub:
		resp, err = srv.handleMailboxSubscribe(ctx, req.AccountPk, call.MboxSub)
		rcname = "MboxSub"
	case *pb.Request_MboxStatus:
		resp, err = srv.handleMailboxStatus(ctx, req.AccountPk, call.MboxStatus)
		rcname = "MboxStatus"

	// --- Mail management ---
	case *pb.Request_MailPush:
		// handleMailPush has its own logic for Case A (Local) and Case B (Remote Node)
		resp, err = srv.handleMailPush(ctx, req.AccountPk, call.MailPush, reqPayload)
		rcname = "MailPush"

	case *pb.Request_MailFetch:
		resp, respPayload, err = srv.handleMailFetch(ctx, req.AccountPk, call.MailFetch)
		rcname = "MailFetch"
	case *pb.Request_MailFlags:
		resp, err = srv.handleMailFlags(ctx, req.AccountPk, call.MailFlags)
		rcname = "MailFlags"
	case *pb.Request_MailMove:
		resp, err = srv.handleMailMove(ctx, req.AccountPk, call.MailMove)
		rcname = "MailMove"
	case *pb.Request_MailDel:
		resp, err = srv.handleMailDelete(ctx, req.AccountPk, call.MailDel)
		rcname = "MailDel"
	case *pb.Request_MailExpunge:
		resp, err = srv.handleMailExpunge(ctx, req.AccountPk, call.MailExpunge)
		rcname = "MailExpunge"
	case *pb.Request_MailIdForSeq:
		resp, err = srv.handleMailIdForSeq(ctx, req.AccountPk, call.MailIdForSeq)
		rcname = "MailIdForSeq"
	case *pb.Request_MailSearch:
		resp, err = srv.handleMailSearch(ctx, req.AccountPk, call.MailSearch)
		rcname = "MailSearch"
	case *pb.Request_MailWatch:
		resp, err = srv.handleMailWatch(ctx, req.AccountPk, call.MailWatch)

	// --- Queue management (Local Account only via Policy) ---
	case *pb.Request_QueueInsert:
		resp, err = srv.handleQueueInsert(ctx, req.AccountPk, call.QueueInsert)
		rcname = "QueueInsert"

	default:
		rcname = "Not Implemented"
		resp = &pb.Response{Success: false, Error: "501 Not Implemented"}
	}
	// srv.log.Printf("RPC: %s : Authorized: %t: Publickey: %s :", rcname, isLocalAccount, req.AccountPk)

	if err != nil {
		srv.log.Printf("RPC Error: %s %v", rcname, err)
		resp = &pb.Response{Success: false, Error: err.Error()}
	}
	return resp, respPayload
}