Release 0.1.4

cmd/hayduk/main.go

@@ -17,7 +17,7 @@ import (
 	"github.com/jolovicdev/hayduk/internal/server"
 )
 
-var version = "0.1.3"
+var version = "0.1.4"
 
 func main() {
 	listen := flag.String("listen", "127.0.0.1:0", "host:port to bind")

internal/engine/commands.go

@@ -15,7 +15,9 @@ func (e *Engine) execCommand(ctx context.Context, operator, method string, param
 	if !knownMethod(method) {
 		return nil, &protocol.ErrorBody{Code: protocol.CodeUnknownMethod, Message: "no such method: " + method}
 	}
-	rpc := e.connectedRPC()
+	// the workspace the command starts in rides with its events: switching
+	// workspaces mid-RPC must not re-attribute the result
+	rpc, ws := e.dispatchScope()
 	if rpc == nil {
 		return nil, notConnected()
 	}
@@ -36,6 +38,16 @@ func (e *Engine) execCommand(ctx context.Context, operator, method string, param
 		if err := con.Write(ctx, p.Command); err != nil {
 			return nil, mapErr(err)
 		}
+		// Readiness is restored only when a read issued after this write
+		// reports the console idle (consoleRead checks the generation):
+		// the buffered prompt must not go out here, or browsers accept
+		// input while msf still runs the command.
+		e.mu.Lock()
+		if e.console != nil {
+			e.consoleWritePending = true
+			e.consoleWriteGen++
+		}
+		e.mu.Unlock()
 		return nil, nil
 
 	case protocol.MethodConsoleTabs:
@@ -111,7 +123,7 @@ func (e *Engine) execCommand(ctx context.Context, operator, method string, param
 		if eb := requireModuleRef(p.Type, p.Name); eb != nil {
 			return nil, eb
 		}
-		return e.moduleExecute(ctx, rpc, operator, p)
+		return e.moduleExecute(ctx, rpc, operator, ws, p)
 
 	case protocol.MethodSessionAttach:
 		var p protocol.SessionRefParams
@@ -173,7 +185,7 @@ func (e *Engine) execCommand(ctx context.Context, operator, method string, param
 		if !validPort(p.LPORT) {
 			return nil, badParam("lport must be between 1 and 65535")
 		}
-		if eb := e.sessionUpgrade(ctx, operator, p); eb != nil {
+		if eb := e.sessionUpgrade(ctx, operator, ws, p); eb != nil {
 			return nil, eb
 		}
 		return nil, nil
@@ -193,17 +205,18 @@ func (e *Engine) execCommand(ctx context.Context, operator, method string, param
 		if p.Name == "" {
 			return nil, badParam("name is required")
 		}
+		// refreshMu spans the whole transition including the workspace RPC:
+		// a periodic refresh that already read the old workspace must not
+		// commit its stale rows after the clear below, nor interleave its
+		// old-workspace page reads with the new workspace's
+		e.refreshMu.Lock()
+		defer e.refreshMu.Unlock()
 		e.mu.Lock()
 		startGen := e.gen
 		e.mu.Unlock()
 		if err := gomsf.NewDbManager(rpc).SetWorkspace(ctx, p.Name); err != nil {
 			return nil, mapErr(err)
 		}
-		// refreshMu spans the whole transition: a periodic refresh that
-		// already read the old workspace must not commit its stale rows
-		// after the clear below
-		e.refreshMu.Lock()
-		defer e.refreshMu.Unlock()
 		// msf already switched: the connection says so and the old
 		// workspace's data must not survive a refresh that fails halfway.
 		// The generation guard keeps a link that was replaced mid-switch
@@ -254,12 +267,12 @@ func (e *Engine) execCommand(ctx context.Context, operator, method string, param
 		if eb := parseParams(params, &p); eb != nil {
 			return nil, eb
 		}
-		return e.hailMary(ctx, operator, p)
+		return e.hailMary(ctx, operator, ws, p)
 	}
 	return nil, &protocol.ErrorBody{Code: protocol.CodeInternal, Message: "unreachable dispatch for " + method}
 }
 
-func (e *Engine) moduleExecute(ctx context.Context, rpc gomsf.RPCCaller, operator string, p protocol.ModuleExecuteParams) (json.RawMessage, *protocol.ErrorBody) {
+func (e *Engine) moduleExecute(ctx context.Context, rpc gomsf.RPCCaller, operator, ws string, p protocol.ModuleExecuteParams) (json.RawMessage, *protocol.ErrorBody) {
 	if rpc == nil {
 		return nil, notConnected()
 	}
@@ -294,15 +307,15 @@ func (e *Engine) moduleExecute(ctx context.Context, rpc gomsf.RPCCaller, operato
 		res, err = mod.Execute(ctx)
 	}
 	if err != nil {
-		e.eventfOp(operator, protocol.LevelError, "%s/%s failed: %v", p.Type, p.Name, err)
+		e.eventfOpIn(ws, operator, protocol.LevelError, "%s/%s failed: %v", p.Type, p.Name, err)
 		return nil, mapErr(err)
 	}
 	// msf answers job 0 when the module runs inline and finishes at once -
 	// there is no job to watch then, and saying "job 0" just confuses
 	if res.JobID > 0 {
-		e.eventfOp(operator, protocol.LevelSuccess, "%s/%s launched as job %d", p.Type, p.Name, res.JobID)
+		e.eventfOpIn(ws, operator, protocol.LevelSuccess, "%s/%s launched as job %d", p.Type, p.Name, res.JobID)
 	} else {
-		e.eventfOp(operator, protocol.LevelSuccess, "%s/%s ran inline", p.Type, p.Name)
+		e.eventfOpIn(ws, operator, protocol.LevelSuccess, "%s/%s ran inline", p.Type, p.Name)
 	}
 	return mustJSON(protocol.ExecPayload{JobID: res.JobID, UUID: res.UUID}), nil
 }
@@ -348,9 +361,19 @@ func (e *Engine) attach(sid string) *protocol.ErrorBody {
 	}
 	previous := e.interactSID
 	e.interactSID = sid
-	e.interactOut = nil
+	if previous != sid {
+		// attachment is idempotent: reopening the already-attached session
+		// must not erase the transcript every connected operator is reading;
+		// that output is already consumed from the RPC stream and cannot be
+		// polled back
+		e.interactOut = nil
+	}
 	mon := e.monitor
-	e.bus.send(protocol.InteractUpdate(&protocol.InteractState{SID: sid}))
+	// the update carries the buffered transcript: browsers replace their
+	// local copy on every interact update and would blank it without this
+	e.bus.send(protocol.InteractUpdate(&protocol.InteractState{
+		SID: sid, Output: string(e.interactOut),
+	}))
 	e.mu.Unlock()
 	if mon != nil {
 		if previous != "" && previous != sid {
@@ -400,7 +423,7 @@ func (e *Engine) sessionWrite(ctx context.Context, p protocol.SessionWriteParams
 	return mapErr(gomsf.NewShellSession(rpc, p.SID).Write(ctx, data))
 }
 
-func (e *Engine) sessionUpgrade(ctx context.Context, operator string, p protocol.SessionUpgradeParams) *protocol.ErrorBody {
+func (e *Engine) sessionUpgrade(ctx context.Context, operator, ws string, p protocol.SessionUpgradeParams) *protocol.ErrorBody {
 	e.mu.Lock()
 	rpc := e.rpc
 	session := e.sessions[p.SID]
@@ -417,7 +440,7 @@ func (e *Engine) sessionUpgrade(ctx context.Context, operator string, p protocol
 	if err := gomsf.NewShellSession(rpc, p.SID).Upgrade(ctx, p.LHOST, p.LPORT); err != nil {
 		return mapErr(err)
 	}
-	e.eventfOp(operator, protocol.LevelSuccess, "session %s upgrading to meterpreter via %s:%d", p.SID, p.LHOST, p.LPORT)
+	e.eventfOpIn(ws, operator, protocol.LevelSuccess, "session %s upgrading to meterpreter via %s:%d", p.SID, p.LHOST, p.LPORT)
 	return nil
 }
 

internal/engine/commands_test.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -18,7 +19,7 @@ func TestModuleExecuteAfterDisconnectIsAnErrorNotACrash(t *testing.T) {
 	e := connectedEngine(t, stdFake())
 	e.Disconnect()
 
-	_, eb := e.moduleExecute(context.Background(), e.connectedRPC(), "dana",
+	_, eb := e.moduleExecute(context.Background(), e.connectedRPC(), "dana", "",
 		protocol.ModuleExecuteParams{Type: "exploit", Name: "windows/smb/a"})
 	if eb == nil || eb.Code != protocol.CodeNotConnected {
 		t.Fatalf("want not_connected after disconnect, got %+v", eb)
@@ -202,3 +203,104 @@ func TestWorkspaceSetSerializesWithInFlightRefresh(t *testing.T) {
 		t.Fatalf("stale rows from the old workspace survived the switch: %d hosts", len(st.Hosts))
 	}
 }
+
+// Reattaching the open session keeps the transcript, and the broadcast
+// carries it so browsers keep their copy too.
+func TestReattachPreservesTranscript(t *testing.T) {
+	e := testEngine(t)
+	e.sessions["1"] = &protocol.SessionState{ID: "1", Type: "shell"}
+	if err := e.attach("1"); err != nil {
+		t.Fatal(err)
+	}
+	e.sessionOutput(nil, gomsf.Event{SessionID: "1", Data: "shared transcript\n"})
+	sub := e.Subscribe()
+	defer sub.Stop()
+	if err := e.attach("1"); err != nil {
+		t.Fatal(err)
+	}
+	if got := e.State().Interact.Output; got != "shared transcript\n" {
+		t.Fatalf("transcript after reattach: %q", got)
+	}
+	select {
+	case m := <-sub.C():
+		up, ok := m.(protocol.ResourceUpdate)
+		if !ok || up.Resource != protocol.ResInteract || up.Interact == nil {
+			t.Fatalf("reattach broadcast %+v", m)
+		}
+		if up.Interact.Output != "shared transcript\n" {
+			t.Fatalf("reattach broadcast carried %q", up.Interact.Output)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("reattach never broadcast")
+	}
+}
+
+// workspace.set must not touch the remote workspace while a refresh is
+// mid-fetch, and the state after the switch must be the new workspace's.
+func TestWorkspaceSwitchWaitsForInFlightRefresh(t *testing.T) {
+	e := testEngine(t)
+	f := stdFake()
+	e.rpc = f
+	e.conn.Workspace = "alpha"
+
+	var ws atomic.Value
+	ws.Store("alpha")
+	setWorkspaceEntered := make(chan struct{})
+	var setOnce sync.Once
+	f.set(gomsf.DbSetWorkspace, func(args ...interface{}) (interface{}, error) {
+		ws.Store(args[0].(string))
+		setOnce.Do(func() { close(setWorkspaceEntered) })
+		return map[string]interface{}{"result": "success"}, nil
+	})
+	f.set(gomsf.DbCurrentWorkspace, func(...interface{}) (interface{}, error) {
+		return map[string]interface{}{"workspace": ws.Load().(string)}, nil
+	})
+	hostRead := make(chan struct{}, 1)
+	release := make(chan struct{})
+	var once sync.Once
+	f.set(gomsf.DbHosts, func(args ...interface{}) (interface{}, error) {
+		once.Do(func() { hostRead <- struct{}{} })
+		<-release
+		return map[string]interface{}{"hosts": []interface{}{
+			map[string]interface{}{"address": args[0].(map[string]interface{})["workspace"].(string) + "-host"},
+		}}, nil
+	})
+	f.set(gomsf.DbServices, func(args ...interface{}) (interface{}, error) {
+		return map[string]interface{}{"services": []interface{}{
+			map[string]interface{}{"host": args[0].(map[string]interface{})["workspace"].(string) + "-host", "port": 80},
+		}}, nil
+	})
+
+	refreshDone := make(chan error, 1)
+	go func() { refreshDone <- e.refreshDB(context.Background()) }()
+	<-hostRead
+
+	switchDone := make(chan struct{})
+	go func() {
+		_, _ = e.Exec(context.Background(), "", protocol.MethodWorkspaceSet, json.RawMessage(`{"name":"beta"}`))
+		close(switchDone)
+	}()
+
+	select {
+	case <-setWorkspaceEntered:
+		t.Fatal("workspace.set ran during an in-flight refresh")
+	case <-time.After(50 * time.Millisecond):
+	}
+
+	close(release)
+	if err := <-refreshDone; err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case <-switchDone:
+	case <-time.After(5 * time.Second):
+		t.Fatal("workspace.set never finished")
+	}
+	st := e.State()
+	if st.Connection.Workspace != "beta" {
+		t.Fatalf("workspace %q", st.Connection.Workspace)
+	}
+	if len(st.Hosts) != 1 || st.Hosts[0].Address != "beta-host" {
+		t.Fatalf("hosts after switch: %+v", st.Hosts)
+	}
+}