diff --git a/cmd/enclave-server/main.go b/cmd/enclave-server/main.go index 1abe9d4..ae74071 100644 --- a/cmd/enclave-server/main.go +++ b/cmd/enclave-server/main.go @@ -25,10 +25,17 @@ import ( const ( vsockPort = 5000 - readDeadline = 30 * time.Second defaultMaxWorkersEnv = "ENCLAVE_MAX_WORKERS" ) +// Connection deadlines. The write deadline starts once the request is handled, +// so a peer that stops reading releases its worker slot. Variables so tests can +// shorten them. +var ( + readDeadline = 30 * time.Second + writeDeadline = 30 * time.Second +) + func main() { km, err := enclave.NewKeyManager() if err != nil { @@ -96,7 +103,10 @@ func handleConnection(conn net.Conn, km *enclave.KeyManager) { } }() - _ = conn.SetReadDeadline(time.Now().Add(readDeadline)) + if err := conn.SetReadDeadline(time.Now().Add(readDeadline)); err != nil { + log.Printf("ERROR: set read deadline: %v", err) + return + } var buf bytes.Buffer if _, err := io.Copy(&buf, conn); err != nil { @@ -114,6 +124,10 @@ func handleConnection(conn net.Conn, km *enclave.KeyManager) { log.Printf("INFO: received request type: %s", base.Type) response := dispatch(base.Type, buf.Bytes(), km) + if err := conn.SetWriteDeadline(time.Now().Add(writeDeadline)); err != nil { + log.Printf("ERROR: set write deadline: %v", err) + return + } if err := json.NewEncoder(conn).Encode(response); err != nil { log.Printf("ERROR: encode response: %v", err) } else { diff --git a/cmd/enclave-server/main_test.go b/cmd/enclave-server/main_test.go new file mode 100644 index 0000000..f219787 --- /dev/null +++ b/cmd/enclave-server/main_test.go @@ -0,0 +1,72 @@ +package main + +import ( + "errors" + "io" + "net" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +// requestConn serves a fixed request on Read and passes writes and deadlines +// through to the embedded conn. +type requestConn struct { + net.Conn + req io.Reader +} + +func (c *requestConn) Read(p []byte) (int, error) { return c.req.Read(p) } + +// readDeadlineFailConn fails SetReadDeadline and records any Read. +type readDeadlineFailConn struct { + net.Conn + read bool +} + +func (*readDeadlineFailConn) SetReadDeadline(time.Time) error { + return errors.New("set read deadline failed") +} + +func (c *readDeadlineFailConn) Read([]byte) (int, error) { + c.read = true + return 0, io.EOF +} + +func TestHandleConnection_WriteDeadlineReleasesStalledPeer(t *testing.T) { + orig := writeDeadline + writeDeadline = 50 * time.Millisecond + t.Cleanup(func() { writeDeadline = orig }) + + // net.Pipe is unbuffered, so the response write blocks until the peer reads. + // The peer never reads. + server, peer := net.Pipe() + t.Cleanup(func() { _ = peer.Close() }) + conn := &requestConn{Conn: server, req: strings.NewReader(`{"type":"ping"}`)} + + done := make(chan struct{}) + go func() { + handleConnection(conn, nil) + close(done) + }() + + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("handleConnection still blocked writing to a peer that never reads") + } +} + +func TestHandleConnection_ClosesWithoutReadingWhenReadDeadlineFails(t *testing.T) { + server, peer := net.Pipe() + t.Cleanup(func() { _ = peer.Close() }) + conn := &readDeadlineFailConn{Conn: server} + + handleConnection(conn, nil) + + assert.False(t, conn.read) + _, err := server.Write([]byte("x")) + assert.Error(t, err) +}