cmd/mbf: handle flags in serve

This enables easier expansion of the protocol.

Signed-off-by: Ophestra <cat@gensokyo.uk>

2 files changed, 23 insertions, 20 deletions

diff --git a/cmd/mbf/daemon.go b/cmd/mbf/daemon.go
index d64643d2..4a5f7d6d 100644
--- a/cmd/mbf/daemon.go
+++ b/cmd/mbf/daemon.go
@@ -35,21 +35,10 @@ const (
 
 // cureFromIR services an IR curing request.
 func cureFromIR(
-	ctx context.Context,
 	cache *pkg.Cache,
 	conn net.Conn,
+	flags uint64,
 ) (pkg.Artifact, error) {
-	go func() { <-ctx.Done(); _ = conn.SetDeadline(time.Now()) }()
-	if err := conn.SetReadDeadline(daemonDeadline()); err != nil {
-		return nil, errors.Join(err, conn.Close())
-	}
-
-	var flagsWire [8]byte
-	if _, err := io.ReadFull(conn, flagsWire[:]); err != nil {
-		return nil, errors.Join(err, conn.Close())
-	}
-	flags := binary.LittleEndian.Uint64(flagsWire[:])
-
 	a, decodeErr := cache.NewDecoder(conn).Decode()
 	if decodeErr != nil {
 		_, err := conn.Write([]byte("\x00" + decodeErr.Error()))
@@ -105,7 +94,26 @@ func serve(
 				continue
 			}
 			wg.Go(func() {
-				if a, _err := cureFromIR(ctx, cm.c, conn); _err != nil {
+				go func() { <-ctx.Done(); _ = conn.SetDeadline(time.Now()) }()
+				if _err := conn.SetReadDeadline(daemonDeadline()); _err != nil {
+					log.Println(_err)
+					if _err = conn.Close(); _err != nil {
+						log.Println(_err)
+					}
+					return
+				}
+
+				var flagsWire [8]byte
+				if _, _err := io.ReadFull(conn, flagsWire[:]); _err != nil {
+					log.Println(_err)
+					if _err = conn.Close(); _err != nil {
+						log.Println(_err)
+					}
+					return
+				}
+				flags := binary.LittleEndian.Uint64(flagsWire[:])
+
+				if a, _err := cureFromIR(cm.c, conn, flags); _err != nil {
 					log.Println(_err)
 				} else {
 					log.Printf(
diff --git a/cmd/mbf/daemon_test.go b/cmd/mbf/daemon_test.go
index 2ec2e76a..6fd95a96 100644
--- a/cmd/mbf/daemon_test.go
+++ b/cmd/mbf/daemon_test.go
@@ -3,7 +3,6 @@ package main
 import (
 	"bytes"
 	"context"
-	"encoding/binary"
 	"errors"
 	"io"
 	"log"
@@ -46,11 +45,7 @@ func TestNoReply(t *testing.T) {
 			}
 		}()
 
-		if _, _err := client.Write(
-			binary.LittleEndian.AppendUint64(nil, remoteNoReply),
-		); _err != nil {
-			panic(_err)
-		} else if _err = c.EncodeAll(
+		if _err := c.EncodeAll(
 			client,
 			pkg.NewFile("check", []byte{0}),
 		); _err != nil {
@@ -60,7 +55,7 @@ func TestNoReply(t *testing.T) {
 		}
 	}()
 
-	a, cureErr := cureFromIR(t.Context(), c, server)
+	a, cureErr := cureFromIR(c, server, remoteNoReply)
 	if cureErr != nil {
 		t.Fatalf("cureFromIR: error = %v", cureErr)
 	}