Add tool to save remote trace sessions

This tool allows testing of trace servers without the need to setup runsc and workloads during testing. More details in the readme file. Updates #4805 PiperOrigin-RevId: 454641010
2026-05-22 17:12:49 -07:00 · 2022-06-13 10:24:06 -07:00
parent 605841baad
commit ab4f6830bc
16 changed files with 997 additions and 183 deletions
@@ -115,7 +115,7 @@ func TestBasic(t *testing.T) {
 	}
 	defer server.Close()

-	endpoint, err := setup(server.Path)
+	endpoint, err := setup(server.Endpoint)
 	if err != nil {
 		t.Fatalf("setup(): %v", err)
 	}
@@ -163,7 +163,7 @@ func TestVersionUnsupported(t *testing.T) {

 	server.SetVersion(0)

-	_, err = setup(server.Path)
+	_, err = setup(server.Endpoint)
 	if err == nil || !strings.Contains(err.Error(), "remote version") {
 		t.Fatalf("Wrong error: %v", err)
 	}
@@ -178,7 +178,7 @@ func TestVersionNewer(t *testing.T) {

 	server.SetVersion(wire.CurrentVersion + 10)

-	endpoint, err := setup(server.Path)
+	endpoint, err := setup(server.Endpoint)
 	if err != nil {
 		t.Fatalf("setup(): %v", err)
 	}
@@ -0,0 +1,19 @@
+load("//tools:defs.bzl", "go_library")
+
+package(licenses = ["notice"])
+
+go_library(
+    name = "server",
+    srcs = ["server.go"],
+    visibility = ["//:sandbox"],
+    deps = [
+        "//pkg/cleanup",
+        "//pkg/log",
+        "//pkg/sentry/seccheck/checkers/remote/wire",
+        "//pkg/sentry/seccheck/points:points_go_proto",
+        "//pkg/sync",
+        "//pkg/unet",
+        "@org_golang_google_protobuf//proto:go_default_library",
+        "@org_golang_x_sys//unix:go_default_library",
+    ],
+)
@@ -0,0 +1,245 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package server provides a common server implementation that can connect with
+// remote.Remote.
+package server
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"os"
+
+	"golang.org/x/sys/unix"
+	"google.golang.org/protobuf/proto"
+	"gvisor.dev/gvisor/pkg/cleanup"
+	"gvisor.dev/gvisor/pkg/log"
+	"gvisor.dev/gvisor/pkg/sentry/seccheck/checkers/remote/wire"
+	pb "gvisor.dev/gvisor/pkg/sentry/seccheck/points/points_go_proto"
+	"gvisor.dev/gvisor/pkg/sync"
+	"gvisor.dev/gvisor/pkg/unet"
+)
+
+// ClientHandler is used to interface with client that connect to the server.
+type ClientHandler interface {
+	// NewClient is called when a new client connects to the server. It returns
+	// a handler that will be bound to the client.
+	NewClient() (MessageHandler, error)
+}
+
+// MessageHandler is used to process messages from a client.
+type MessageHandler interface {
+	// Message processes a single message. raw contains the entire unparsed
+	// message. hdr is the parser message header and payload is the unparsed
+	// message data.
+	Message(raw []byte, hdr wire.Header, payload []byte) error
+
+	// Version returns what wire version of the protocol is supported.
+	Version() uint32
+
+	// Close closes the handler.
+	Close()
+}
+
+type client struct {
+	socket  *unet.Socket
+	handler MessageHandler
+}
+
+func (c client) close() {
+	_ = c.socket.Close()
+	c.handler.Close()
+}
+
+// CommonServer provides common functionality to connect and process messages
+// from different clients. Implementors decide how clients and messages are
+// handled, e.g. counting messages for testing.
+type CommonServer struct {
+	// Endpoint is the path to the socket that the server listens to.
+	Endpoint string
+
+	socket *unet.ServerSocket
+
+	handler ClientHandler
+
+	cond sync.Cond
+
+	// +checklocks:cond.L
+	clients []client
+}
+
+// Init initializes the server. It must be called before it is used.
+func (s *CommonServer) Init(path string, handler ClientHandler) {
+	s.Endpoint = path
+	s.handler = handler
+	s.cond = sync.Cond{L: &sync.Mutex{}}
+}
+
+// Start creates the socket file and listens for new connections.
+func (s *CommonServer) Start() error {
+	socket, err := unix.Socket(unix.AF_UNIX, unix.SOCK_SEQPACKET, 0)
+	if err != nil {
+		return fmt.Errorf("socket(AF_UNIX, SOCK_SEQPACKET, 0): %w", err)
+	}
+	cu := cleanup.Make(func() {
+		_ = unix.Close(socket)
+	})
+	defer cu.Clean()
+
+	sa := &unix.SockaddrUnix{Name: s.Endpoint}
+	if err := unix.Bind(socket, sa); err != nil {
+		return fmt.Errorf("bind(%q): %w", s.Endpoint, err)
+	}
+
+	s.socket, err = unet.NewServerSocket(socket)
+	if err != nil {
+		return err
+	}
+	cu.Add(func() { s.socket.Close() })
+
+	if err := s.socket.Listen(); err != nil {
+		return err
+	}
+
+	go s.run()
+	cu.Release()
+	return nil
+}
+
+func (s *CommonServer) run() {
+	for {
+		socket, err := s.socket.Accept()
+		if err != nil {
+			// EBADF returns when the socket closes.
+			if !errors.Is(err, unix.EBADF) {
+				log.Warningf("socket.Accept(): %v", err)
+			}
+			return
+		}
+		msgHandler, err := s.handler.NewClient()
+		if err != nil {
+			log.Warningf("handler.NewClient: %v", err)
+			return
+		}
+		client := client{
+			socket:  socket,
+			handler: msgHandler,
+		}
+		s.cond.L.Lock()
+		s.clients = append(s.clients, client)
+		s.cond.Broadcast()
+		s.cond.L.Unlock()
+
+		if err := s.handshake(client); err != nil {
+			log.Warningf(err.Error())
+			s.closeClient(client)
+			continue
+		}
+		go s.handleClient(client)
+	}
+}
+
+// handshake performs version exchange with client. See common.proto for details
+// about the protocol.
+func (s *CommonServer) handshake(client client) error {
+	var in [1024]byte
+	read, err := client.socket.Read(in[:])
+	if err != nil {
+		return fmt.Errorf("reading handshake message: %w", err)
+	}
+	hsIn := pb.Handshake{}
+	if err := proto.Unmarshal(in[:read], &hsIn); err != nil {
+		return fmt.Errorf("unmarshalling handshake message: %w", err)
+	}
+	if hsIn.Version != wire.CurrentVersion {
+		return fmt.Errorf("wrong version number, want: %d, got, %d", wire.CurrentVersion, hsIn.Version)
+	}
+
+	hsOut := pb.Handshake{Version: client.handler.Version()}
+	out, err := proto.Marshal(&hsOut)
+	if err != nil {
+		return fmt.Errorf("marshalling handshake message: %w", err)
+	}
+	if _, err := client.socket.Write(out); err != nil {
+		return fmt.Errorf("sending handshake message: %w", err)
+	}
+	return nil
+}
+
+func (s *CommonServer) handleClient(client client) {
+	defer s.closeClient(client)
+
+	var buf = make([]byte, 1024*1024)
+	for {
+		read, err := client.socket.Read(buf)
+		if err != nil {
+			if errors.Is(err, io.EOF) || errors.Is(err, unix.EBADF) {
+				// Both errors indicate that the socket has been closed.
+				return
+			}
+			panic(err)
+		}
+		if read < wire.HeaderStructSize {
+			panic("message too small")
+		}
+		hdr := wire.Header{}
+		hdr.UnmarshalUnsafe(buf[0:wire.HeaderStructSize])
+		if read < int(hdr.HeaderSize) {
+			panic(fmt.Sprintf("message truncated, header size: %d, read: %d", hdr.HeaderSize, read))
+		}
+		if err := client.handler.Message(buf[:read], hdr, buf[hdr.HeaderSize:read]); err != nil {
+			panic(err)
+		}
+	}
+}
+
+func (s *CommonServer) closeClient(client client) {
+	client.close()
+
+	// Stop tracking this client.
+	s.cond.L.Lock()
+	for i, c := range s.clients {
+		if c == client {
+			s.clients = append(s.clients[:i], s.clients[i+1:]...)
+			break
+		}
+	}
+	s.cond.Broadcast()
+	s.cond.L.Unlock()
+}
+
+// Close stops listening and closes all connections.
+func (s *CommonServer) Close() {
+	if s.socket != nil {
+		_ = s.socket.Close()
+	}
+	s.cond.L.Lock()
+	for _, client := range s.clients {
+		client.close()
+	}
+	s.clients = nil
+	s.cond.Broadcast()
+	s.cond.L.Unlock()
+	_ = os.Remove(s.Endpoint)
+}
+
+// WaitForNoClients waits until the number of clients connected reaches 0.
+func (s *CommonServer) WaitForNoClients() {
+	s.cond.L.Lock()
+	defer s.cond.L.Unlock()
+	for len(s.clients) > 0 {
+		s.cond.Wait()
+	}
+}
@@ -8,13 +8,9 @@ go_library(
    srcs = ["server.go"],
    visibility = ["//:sandbox"],
    deps = [
-        "//pkg/cleanup",
-        "//pkg/log",
+        "//pkg/sentry/seccheck/checkers/remote/server",
        "//pkg/sentry/seccheck/checkers/remote/wire",
        "//pkg/sentry/seccheck/points:points_go_proto",
        "//pkg/sync",
-        "//pkg/unet",
-        "@org_golang_google_protobuf//proto:go_default_library",
-        "@org_golang_x_sys//unix:go_default_library",
    ],
 )
@@ -16,33 +16,23 @@
 package test

 import (
-	"errors"
-	"fmt"
 	"io/ioutil"
 	"os"
 	"path/filepath"

-	"golang.org/x/sys/unix"
-	"google.golang.org/protobuf/proto"
-	"gvisor.dev/gvisor/pkg/cleanup"
-	"gvisor.dev/gvisor/pkg/log"
+	"gvisor.dev/gvisor/pkg/sentry/seccheck/checkers/remote/server"
 	"gvisor.dev/gvisor/pkg/sentry/seccheck/checkers/remote/wire"
 	pb "gvisor.dev/gvisor/pkg/sentry/seccheck/points/points_go_proto"
 	"gvisor.dev/gvisor/pkg/sync"
-	"gvisor.dev/gvisor/pkg/unet"
 )

 // Server is the counterpart to the checkers.Remote. It receives connections
 // remote checkers and stores all points that it receives.
 type Server struct {
-	Path   string
-	socket *unet.ServerSocket
+	server.CommonServer

 	cond sync.Cond

-	// +checklocks:cond.L
-	clients []*unet.Socket
-
 	// +checklocks:cond.L
 	points []Message

@@ -67,147 +57,21 @@ func NewServer() (*Server, error) {
 	if err != nil {
 		return nil, err
 	}
-	server, err := newServerPath(filepath.Join(dir, "remote.sock"))
-	if err != nil {
-		_ = os.RemoveAll(dir)
-		return nil, err
-	}
-	return server, nil
-}
-
-func newServerPath(path string) (*Server, error) {
-	socket, err := unix.Socket(unix.AF_UNIX, unix.SOCK_SEQPACKET, 0)
-	if err != nil {
-		return nil, fmt.Errorf("socket(AF_UNIX, SOCK_SEQPACKET, 0): %w", err)
-	}
-	cu := cleanup.Make(func() {
-		_ = unix.Close(socket)
-	})
-	defer cu.Clean()
-
-	sa := &unix.SockaddrUnix{Name: path}
-	if err := unix.Bind(socket, sa); err != nil {
-		return nil, fmt.Errorf("bind(%q): %w", path, err)
-	}
-
-	ss, err := unet.NewServerSocket(socket)
-	if err != nil {
-		return nil, err
-	}
-	cu.Add(func() { ss.Close() })
-
-	if err := ss.Listen(); err != nil {
-		return nil, err
-	}
-
-	server := &Server{
-		Path:    path,
-		socket:  ss,
+	s := &Server{
 		version: wire.CurrentVersion,
 		cond:    sync.Cond{L: &sync.Mutex{}},
 	}
-	go server.run()
-	cu.Release()
-	return server, nil
+	s.CommonServer.Init(filepath.Join(dir, "remote.sock"), s)
+	if err := s.CommonServer.Start(); err != nil {
+		_ = os.RemoveAll(dir)
+		return nil, err
+	}
+	return s, nil
 }

-func (s *Server) run() {
-	for {
-		client, err := s.socket.Accept()
-		if err != nil {
-			// EBADF returns when the socket closes.
-			if !errors.Is(err, unix.EBADF) {
-				log.Warningf("socket.Accept(): %v", err)
-			}
-			return
-		}
-		if err := s.handshake(client); err != nil {
-			log.Warningf(err.Error())
-			_ = client.Close()
-			continue
-		}
-		s.cond.L.Lock()
-		s.clients = append(s.clients, client)
-		s.cond.Broadcast()
-		s.cond.L.Unlock()
-		go s.handleClient(client)
-	}
-}
-
-// handshake performs version exchange with client. See common.proto for details
-// about the protocol.
-func (s *Server) handshake(client *unet.Socket) error {
-	var in [1024]byte
-	read, err := client.Read(in[:])
-	if err != nil {
-		return fmt.Errorf("reading handshake message: %w", err)
-	}
-	hsIn := pb.Handshake{}
-	if err := proto.Unmarshal(in[:read], &hsIn); err != nil {
-		return fmt.Errorf("unmarshalling handshake message: %w", err)
-	}
-	if hsIn.Version != wire.CurrentVersion {
-		return fmt.Errorf("wrong version number, want: %d, got, %d", wire.CurrentVersion, hsIn.Version)
-	}
-
-	s.mu.Lock()
-	v := s.version
-	s.mu.Unlock()
-	hsOut := pb.Handshake{Version: v}
-	out, err := proto.Marshal(&hsOut)
-	if err != nil {
-		return fmt.Errorf("marshalling handshake message: %w", err)
-	}
-	if _, err := client.Write(out); err != nil {
-		return fmt.Errorf("sending handshake message: %w", err)
-	}
-	return nil
-}
-
-func (s *Server) handleClient(client *unet.Socket) {
-	defer func() {
-		s.cond.L.Lock()
-		for i, c := range s.clients {
-			if c == client {
-				s.clients = append(s.clients[:i], s.clients[i+1:]...)
-				break
-			}
-		}
-		s.cond.Broadcast()
-		s.cond.L.Unlock()
-		_ = client.Close()
-	}()
-
-	var buf = make([]byte, 1024*1024)
-	for {
-		read, err := client.Read(buf)
-		if err != nil {
-			return
-		}
-		if read == 0 {
-			return
-		}
-		if read < wire.HeaderStructSize {
-			panic("invalid message")
-		}
-		hdr := wire.Header{}
-		hdr.UnmarshalUnsafe(buf[0:wire.HeaderStructSize])
-		if read < int(hdr.HeaderSize) {
-			panic(fmt.Sprintf("message truncated, header size: %d, readL %d", hdr.HeaderSize, read))
-		}
-
-		msgSize := read - int(hdr.HeaderSize)
-		msg := Message{
-			MsgType: pb.MessageType(hdr.MessageType),
-			Msg:     make([]byte, msgSize),
-		}
-		copy(msg.Msg, buf[hdr.HeaderSize:read])
-
-		s.cond.L.Lock()
-		s.points = append(s.points, msg)
-		s.cond.Broadcast()
-		s.cond.L.Unlock()
-	}
+// NewClient returns a new MessageHandler to process messages.
+func (s *Server) NewClient() (server.MessageHandler, error) {
+	return &msgHandler{owner: s}, nil
 }

 // Count return the number of points it has received.
@@ -236,19 +100,6 @@ func (s *Server) GetPoints() []Message {
 	return cpy
 }

-// Close stops listenning and closes all connections.
-func (s *Server) Close() {
-	_ = s.socket.Close()
-	s.cond.L.Lock()
-	for _, client := range s.clients {
-		_ = client.Close()
-	}
-	s.clients = nil
-	s.cond.Broadcast()
-	s.cond.L.Unlock()
-	_ = os.Remove(s.Path)
-}
-
 // WaitForCount waits for the number of points to reach the desired number.
 func (s *Server) WaitForCount(count int) {
 	s.cond.L.Lock()
@@ -259,18 +110,38 @@ func (s *Server) WaitForCount(count int) {
 	return
 }

-// WaitForNoClients waits until the number of clients connected reaches 0.
-func (s *Server) WaitForNoClients() {
-	s.cond.L.Lock()
-	defer s.cond.L.Unlock()
-	for len(s.clients) > 0 {
-		s.cond.Wait()
-	}
-}
-
 // SetVersion sets the version to be used in handshake.
 func (s *Server) SetVersion(newVersion uint32) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	s.version = newVersion
 }
+
+type msgHandler struct {
+	owner *Server
+}
+
+// Message stores the message type and payload.
+func (m *msgHandler) Message(_ []byte, hdr wire.Header, payload []byte) error {
+	msg := Message{
+		MsgType: pb.MessageType(hdr.MessageType),
+		Msg:     make([]byte, len(payload)),
+	}
+	copy(msg.Msg, payload)
+
+	m.owner.cond.L.Lock()
+	defer m.owner.cond.L.Unlock()
+	m.owner.points = append(m.owner.points, msg)
+	m.owner.cond.Broadcast()
+	return nil
+}
+
+// Version returns the wire version supported or overriden by SetVersion.
+func (m *msgHandler) Version() uint32 {
+	m.owner.mu.Lock()
+	defer m.owner.mu.Unlock()
+	return m.owner.version
+}
+
+// Close implements server.MessageHandler.
+func (m *msgHandler) Close() {}
@@ -69,7 +69,7 @@ func TestTraceStartup(t *testing.T) {
 							ContextFields: []string{"container_id"},
 						},
 					},
-					Sinks: []seccheck.SinkConfig{remoteSinkConfig(server.Path)},
+					Sinks: []seccheck.SinkConfig{remoteSinkConfig(server.Endpoint)},
 				},
 			}
 			encoder := json.NewEncoder(podInitConfig)
@@ -148,7 +148,7 @@ func TestTraceLifecycle(t *testing.T) {
 				ContextFields: []string{"container_id"},
 			},
 		},
-		Sinks: []seccheck.SinkConfig{remoteSinkConfig(server.Path)},
+		Sinks: []seccheck.SinkConfig{remoteSinkConfig(server.Endpoint)},
 	}
 	if err := cont.Sandbox.CreateTraceSession(&session, false); err != nil {
 		t.Fatalf("CreateTraceSession(): %v", err)
@@ -248,7 +248,7 @@ func TestTraceForceCreate(t *testing.T) {
 		Points: []seccheck.PointConfig{
 			{Name: "sentry/exit_notify_parent"},
 		},
-		Sinks: []seccheck.SinkConfig{remoteSinkConfig(server.Path)},
+		Sinks: []seccheck.SinkConfig{remoteSinkConfig(server.Endpoint)},
 	}
 	if err := cont.Sandbox.CreateTraceSession(&session, false); err != nil {
 		t.Fatalf("CreateTraceSession(): %v", err)
@@ -277,7 +277,7 @@ func TestTraceForceCreate(t *testing.T) {
 		Points: []seccheck.PointConfig{
 			{Name: "sentry/task_exit"},
 		},
-		Sinks: []seccheck.SinkConfig{remoteSinkConfig(server.Path)},
+		Sinks: []seccheck.SinkConfig{remoteSinkConfig(server.Endpoint)},
 	}
 	if err := cont.Sandbox.CreateTraceSession(&session, true); err != nil {
 		t.Fatalf("CreateTraceSession(force): %v", err)
@@ -48,7 +48,7 @@ func TestAll(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	cfg, err := buildPodConfig(runsc, server.Path)
+	cfg, err := buildPodConfig(runsc, server.Endpoint)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -0,0 +1,36 @@
+load("//tools:defs.bzl", "go_library", "go_test")
+
+package(licenses = ["notice"])
+
+go_library(
+    name = "tracereplay",
+    srcs = [
+        "replay.go",
+        "save.go",
+        "tracereplay.go",
+    ],
+    visibility = [
+        "//tools/tracereplay:__subpackages__",
+    ],
+    deps = [
+        "//pkg/atomicbitops",
+        "//pkg/log",
+        "//pkg/sentry/seccheck/checkers/remote/server",
+        "//pkg/sentry/seccheck/checkers/remote/wire",
+        "//pkg/sentry/seccheck/points:points_go_proto",
+        "@org_golang_google_protobuf//proto:go_default_library",
+        "@org_golang_x_sys//unix:go_default_library",
+    ],
+)
+
+go_test(
+    name = "tracereplay_test",
+    srcs = ["tracereplay_test.go"],
+    data = [
+        "testdata/client-0001",
+    ],
+    library = ":tracereplay",
+    deps = [
+        "//pkg/test/testutil",
+    ],
+)
@@ -0,0 +1,77 @@
+# What is it?
+
+The `tracereplay` tool can save `runsc trace` sessions to a file, and later
+replay the same sequence of messages. This can be used to run tests that rely on
+the messages without the need to setup runsc, configure trace sessions, and run
+specific workloads.
+
+# How to use it?
+
+The `tracereplay save` command starts a server that listens to new connections
+from runsc and creates a trace file for each runsc instance that connects to it.
+The command below starts a server on listening on `/tmp/gvisor_events.sock` and
+writes trace files to `/tmp/trace` directory:
+
+```shell
+$ tracereplay save --endpoint=/tmp/gvisor_events.sock --out=/tmp/trace
+```
+
+When you execute runsc configured with a trace session using the `remote` sink
+connecting to `/tmp/gvisor_events.sock`, all messages will be saved to a file
+under `/tmp/trace`. For example, if you run the following commands, runsc will
+connect to the server above and all trace points triggered by the workload will
+be stored in the save file:
+
+```shell
+$ cat > /tmp/pod_init.json <<EOL
+{
+  "trace_session": {
+    "name": "Default",
+    "points": [
+      {
+        "name": "container/start"
+      }
+    ],
+    "sinks": [
+      {
+        "name": "remote",
+        "config": {
+          "endpoint": "/tmp/gvisor_events.sock"
+        }
+      }
+    ]
+  }
+}
+EOL
+$ runsc --rootless --network=none --pod-init-config=/tmp/pod_init.json do /bin/true
+```
+
+You should see the following output from `tracereplay save`:
+
+```
+New client connected, writing to: "/tmp/trace/client-0001"
+Closing client, wrote 1 messages to "/tmp/trace/client-0001"
+```
+
+You can then use the `tracereplay replay` command to replay the exact same
+messages anytime and as many times as you want. Here is an example using the
+file created above:
+
+```shell
+$ tracereplay replay --endpoint=/tmp/gvisor_events.sock --in=/tmp/trace/client-0001
+Handshake completed
+Replaying message: 1
+Done
+```
+
+If you want to see the messages that are stored in the file, you can setup the
+example server provided in `examples/seccheck:server_cc` and replay the save
+file using the same command above. Here is the output you would get:
+
+```shell
+$ bazel run examples/seccheck:server_cc
+Socket address /tmp/gvisor_events.sock
+Connection accepted
+Start => id:     "runsc-865139" cwd: "/home/fvoznika" args: "/bin/true"
+Connection closed
+```
@@ -0,0 +1,15 @@
+load("//tools:defs.bzl", "go_binary")
+
+package(licenses = ["notice"])
+
+go_binary(
+    name = "tracereplay",
+    srcs = [
+        "main.go",
+    ],
+    deps = [
+        "//runsc/flag",
+        "//tools/tracereplay",
+        "@com_github_google_subcommands//:go_default_library",
+    ],
+)
@@ -0,0 +1,152 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package main implements a tool that can save and replay messages from
+// issued from remote.Remote.
+package main
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"os/signal"
+
+	"github.com/google/subcommands"
+	"gvisor.dev/gvisor/runsc/flag"
+	"gvisor.dev/gvisor/tools/tracereplay"
+)
+
+func main() {
+	subcommands.Register(subcommands.HelpCommand(), "")
+	subcommands.Register(subcommands.FlagsCommand(), "")
+	subcommands.Register(&saveCmd{}, "")
+	subcommands.Register(&replayCmd{}, "")
+	flag.CommandLine.Parse(os.Args[1:])
+	os.Exit(int(subcommands.Execute(context.Background())))
+}
+
+// saveCmd implements subcommands.Command for the "save" command.
+type saveCmd struct {
+	endpoint string
+	out      string
+	prefix   string
+}
+
+// Name implements subcommands.Command.
+func (*saveCmd) Name() string {
+	return "save"
+}
+
+// Synopsis implements subcommands.Command.
+func (*saveCmd) Synopsis() string {
+	return "save trace sessions to files"
+}
+
+// Usage implements subcommands.Command.
+func (*saveCmd) Usage() string {
+	return `save [flags] - save trace sessions to files
+`
+}
+
+// SetFlags implements subcommands.Command.
+func (c *saveCmd) SetFlags(f *flag.FlagSet) {
+	f.StringVar(&c.endpoint, "endpoint", "", "path to trace server endpoint to connect")
+	f.StringVar(&c.out, "out", "./replay", "path to a directory where trace files will be saved")
+	f.StringVar(&c.prefix, "prefix", "client-", "name to be prefixed to each trace file")
+}
+
+// Execute implements subcommands.Command.
+func (c *saveCmd) Execute(_ context.Context, f *flag.FlagSet, args ...interface{}) subcommands.ExitStatus {
+	if f.NArg() > 0 {
+		fmt.Fprintf(os.Stderr, "unexpected argument: %s\n", f.Args())
+		return subcommands.ExitUsageError
+	}
+	if len(c.endpoint) == 0 {
+		fmt.Fprintf(os.Stderr, "--endpoint is required\n")
+		return subcommands.ExitUsageError
+	}
+	_ = os.Remove(c.endpoint)
+
+	server := tracereplay.NewSave(c.endpoint, c.out, c.prefix)
+	defer server.Close()
+
+	if err := server.Start(); err != nil {
+		fmt.Fprintf(os.Stderr, "starting server: %v\n", err)
+		return subcommands.ExitFailure
+	}
+
+	ch := make(chan os.Signal)
+	signal.Notify(ch, os.Interrupt)
+
+	done := make(chan struct{})
+	go func() {
+		<-ch
+		fmt.Printf("Ctrl-C pressed, stopping.\n")
+		done <- struct{}{}
+	}()
+
+	fmt.Printf("Listening on %q. Press ctrl-C to stop...\n", c.endpoint)
+	<-done
+	return subcommands.ExitSuccess
+}
+
+// replayCmd implements subcommands.Command for the "replay" command.
+type replayCmd struct {
+	endpoint string
+	in       string
+}
+
+// Name implements subcommands.Command.
+func (*replayCmd) Name() string {
+	return "replay"
+}
+
+// Synopsis implements subcommands.Command.
+func (*replayCmd) Synopsis() string {
+	return "replay a trace session from a file"
+}
+
+// Usage implements subcommands.Command.
+func (*replayCmd) Usage() string {
+	return `replay [flags] - replay a trace session from a file
+`
+}
+
+// SetFlags implements subcommands.Command.
+func (c *replayCmd) SetFlags(f *flag.FlagSet) {
+	f.StringVar(&c.endpoint, "endpoint", "", "path to trace server endpoint to connect")
+	f.StringVar(&c.in, "in", "", "path to trace file containing messages to be replayed")
+}
+
+// Execute implements subcommands.Command.
+func (c *replayCmd) Execute(_ context.Context, f *flag.FlagSet, _ ...interface{}) subcommands.ExitStatus {
+	if f.NArg() > 0 {
+		fmt.Fprintf(os.Stderr, "unexpected argument: %s\n", f.Args())
+		return subcommands.ExitUsageError
+	}
+	if len(c.in) == 0 {
+		fmt.Fprintf(os.Stderr, "--in is required\n")
+		return subcommands.ExitUsageError
+	}
+
+	r := tracereplay.Replay{
+		Endpoint: c.endpoint,
+		In:       c.in,
+	}
+	if err := r.Execute(); err != nil {
+		fmt.Fprintln(os.Stderr, err)
+		return subcommands.ExitFailure
+	}
+	return subcommands.ExitSuccess
+}
@@ -0,0 +1,133 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tracereplay
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"os"
+
+	"golang.org/x/sys/unix"
+	"google.golang.org/protobuf/proto"
+	"gvisor.dev/gvisor/pkg/log"
+	pb "gvisor.dev/gvisor/pkg/sentry/seccheck/points/points_go_proto"
+)
+
+// Replay implements the functionality required for the "replay" command.
+type Replay struct {
+	Endpoint string
+	In       string
+}
+
+// Execute connects to the remote endpoint and replays all messages stored in
+// the `In` file.
+func (r *Replay) Execute() error {
+	socket, err := connect(r.Endpoint)
+	if err != nil {
+		return err
+	}
+	defer socket.Close()
+
+	f, err := os.Open(r.In)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	hdr := make([]byte, len(signature))
+	if err := readFull(f, hdr); err != nil {
+		return err
+	}
+	if string(hdr) != signature {
+		return fmt.Errorf("%q is not a replay file", r.In)
+	}
+
+	cfgJSON, err := readWithSize(f)
+	if err != nil {
+		return err
+	}
+	cfg := Config{}
+	if err := json.Unmarshal(cfgJSON, &cfg); err != nil {
+		return err
+	}
+	if err := handshake(socket, cfg.Version); err != nil {
+		return err
+	}
+	fmt.Printf("Handshake completed\n")
+
+	for count := 1; ; count++ {
+		bytes, err := readWithSize(f)
+		if err != nil {
+			if errors.Is(err, io.EOF) {
+				break
+			}
+			return err
+		}
+		fmt.Printf("\rReplaying message: %d", count)
+		if _, err := socket.Write(bytes); err != nil {
+			return err
+		}
+	}
+	fmt.Printf("\nDone\n")
+	return nil
+}
+
+func connect(endpoint string) (*os.File, error) {
+	log.Debugf("Connecting to %q", endpoint)
+	socket, err := unix.Socket(unix.AF_UNIX, unix.SOCK_SEQPACKET, 0)
+	if err != nil {
+		return nil, fmt.Errorf("socket(AF_UNIX, SOCK_SEQPACKET, 0): %w", err)
+	}
+	f := os.NewFile(uintptr(socket), endpoint)
+
+	addr := unix.SockaddrUnix{Name: endpoint}
+	if err := unix.Connect(int(f.Fd()), &addr); err != nil {
+		_ = f.Close()
+		return nil, fmt.Errorf("connect(%q): %w", endpoint, err)
+	}
+	return f, nil
+}
+
+// See common.proto for details about the handshake protocol.
+func handshake(socket *os.File, version uint32) error {
+	hsOut := pb.Handshake{Version: version}
+	out, err := proto.Marshal(&hsOut)
+	if err != nil {
+		return err
+	}
+	if _, err := socket.Write(out); err != nil {
+		return fmt.Errorf("sending handshake message: %w", err)
+	}
+
+	in := make([]byte, 10240)
+	read, err := socket.Read(in)
+	if err != nil && !errors.Is(err, io.EOF) {
+		return fmt.Errorf("reading handshake message: %w", err)
+	}
+	// Protect against the handshake becoming larger than the buffer.
+	if read == len(in) {
+		return fmt.Errorf("handshake message too big")
+	}
+	hsIn := pb.Handshake{}
+	if err := proto.Unmarshal(in[:read], &hsIn); err != nil {
+		return fmt.Errorf("unmarshalling handshake message: %w", err)
+	}
+
+	// Just validate that the message can unmarshall and accept any version from
+	// the server. Will try to replay and see what happens...
+	return nil
+}
@@ -0,0 +1,111 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tracereplay
+
+import (
+	"encoding/json"
+	"fmt"
+	"os"
+	"path/filepath"
+
+	"gvisor.dev/gvisor/pkg/atomicbitops"
+	"gvisor.dev/gvisor/pkg/sentry/seccheck/checkers/remote/server"
+	"gvisor.dev/gvisor/pkg/sentry/seccheck/checkers/remote/wire"
+)
+
+// Save implements the functionality required for the "save" command.
+type Save struct {
+	server.CommonServer
+	dir         string
+	prefix      string
+	clientCount atomicbitops.Uint64
+}
+
+var _ server.ClientHandler = (*Save)(nil)
+
+// NewSave creates a new Save instance.
+func NewSave(endpoint, dir, prefix string) *Save {
+	s := &Save{dir: dir, prefix: prefix}
+	s.CommonServer.Init(endpoint, s)
+	return s
+}
+
+// Start starts the server.
+func (s *Save) Start() error {
+	if err := os.MkdirAll(s.dir, 0755); err != nil {
+		return err
+	}
+	return s.CommonServer.Start()
+}
+
+// NewClient creates a new file for the client and writes messages to it.
+//
+// The file format starts with a string signature to make it easy to check that
+// it's a trace file. The signature is followed by a JSON configuration that
+// contains information required to process the file. Next, there are a sequence
+// of messages. Both JSON and messages are prefixed by an uint64 with their
+// size.
+//
+// Ex:
+//   signature <size>Config JSON [<size>message]*
+func (s *Save) NewClient() (server.MessageHandler, error) {
+	seq := s.clientCount.Add(1)
+	filename := filepath.Join(s.dir, fmt.Sprintf("%s%04d", s.prefix, seq))
+	fmt.Printf("New client connected, writing to: %q\n", filename)
+
+	out, err := os.Create(filename)
+	if err != nil {
+		return nil, err
+	}
+	if _, err := out.Write([]byte(signature)); err != nil {
+		return nil, err
+	}
+
+	handler := &msgHandler{out: out}
+
+	cfg, err := json.Marshal(Config{Version: handler.Version()})
+	if err != nil {
+		return nil, err
+	}
+	if err := writeWithSize(out, cfg); err != nil {
+		return nil, err
+	}
+
+	return handler, nil
+}
+
+type msgHandler struct {
+	out          *os.File
+	messageCount atomicbitops.Uint64
+}
+
+var _ server.MessageHandler = (*msgHandler)(nil)
+
+// Version implements server.MessageHandler.
+func (m *msgHandler) Version() uint32 {
+	return wire.CurrentVersion
+}
+
+// Message saves the message to the client file.
+func (m *msgHandler) Message(raw []byte, _ wire.Header, _ []byte) error {
+	m.messageCount.Add(1)
+	return writeWithSize(m.out, raw)
+}
+
+// Close closes the client file.
+func (m *msgHandler) Close() {
+	fmt.Printf("Closing client, wrote %d messages to %q\n", m.messageCount.Load(), m.out.Name())
+	_ = m.out.Close()
+}
@@ -0,0 +1,83 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package tracereplay implements a tool that can save and replay messages
+// issued from remote.Remote.
+package tracereplay
+
+import (
+	"encoding/binary"
+	"fmt"
+	"io"
+	"os"
+)
+
+const signature = "tracereplay file"
+
+func writeSize(w io.Writer, val int) error {
+	var bin [8]byte
+	binary.LittleEndian.PutUint64(bin[:], uint64(val))
+	_, err := w.Write(bin[:])
+	return err
+}
+
+func readSize(r io.Reader) (int, error) {
+	var bin [8]byte
+	if read, err := r.Read(bin[:]); err != nil {
+		return 0, err
+	} else if read != 8 {
+		return 0, fmt.Errorf("truncated read (%d bytes)", read)
+	}
+	size := int(binary.LittleEndian.Uint64(bin[:]))
+	// Prevent returning a too large size to avoid OOMs.
+	if size > 1024*1024 {
+		return 0, fmt.Errorf("size is too big: %d", size)
+	}
+	return size, nil
+}
+
+func writeWithSize(f *os.File, buf []byte) error {
+	if err := writeSize(f, len(buf)); err != nil {
+		return err
+	}
+	_, err := f.Write(buf)
+	return err
+}
+
+func readWithSize(r io.Reader) ([]byte, error) {
+	size, err := readSize(r)
+	if err != nil {
+		return nil, err
+	}
+	bytes := make([]byte, size)
+	if err := readFull(r, bytes); err != nil {
+		return nil, err
+	}
+	return bytes, nil
+}
+
+func readFull(r io.Reader, dest []byte) error {
+	if read, err := r.Read(dest); err != nil {
+		return err
+	} else if read < len(dest) {
+		return fmt.Errorf("truncated read. Read %d bytes, expected %d bytes", read, len(dest))
+	}
+	return nil
+}
+
+// Config contains information required to replay messages from a file.
+type Config struct {
+	// Version is the wire format saved in the file.
+	Version uint32 `json:"version"`
+}
@@ -0,0 +1,76 @@
+// Copyright 2022 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tracereplay
+
+import (
+	"bytes"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"gvisor.dev/gvisor/pkg/test/testutil"
+)
+
+// TestBasic uses a pre-generated file that is replayed into a save process.
+// Then verifies that the generated file looks exactly the same as the original.
+// In other words, it is doing `replay original | save new`, then checking if
+// `original == new`.
+func TestBasic(t *testing.T) {
+	dir, err := os.MkdirTemp(testutil.TmpDir(), "tracereplay")
+	if err != nil {
+		t.Fatal(err)
+	}
+	endpoint := filepath.Join(dir, "tracereplay.sock")
+
+	// Start a new save server to store the replayed file. This tests that save
+	// communicates with clients correctly and generates a valid file.
+	s := NewSave(endpoint, filepath.Join(dir, "out"), "test-")
+	defer s.Close()
+
+	if err := s.Start(); err != nil {
+		t.Fatal(err)
+	}
+
+	// Then replay the re-generated file. This tests that replay can connect to
+	// a server and process the generated file.
+	r := Replay{}
+	r.Endpoint = endpoint
+
+	const testdata = "tools/tracereplay/testdata/client-0001"
+	r.In, err = testutil.FindFile(testdata)
+	if err != nil {
+		t.Fatalf("FindFile(%q): %v", testdata, err)
+	}
+
+	if err := r.Execute(); err != nil {
+		t.Fatal(err)
+	}
+
+	// Wait until all messages are processed and client disconnects.
+	s.WaitForNoClients()
+
+	// The generated file must be an exact copy of the original file.
+	want, err := os.ReadFile(r.In)
+	if err != nil {
+		t.Fatal(err)
+	}
+	got, err := os.ReadFile(filepath.Join(dir, "out", "test-0001"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Equal(want, got) {
+		t.Errorf("files don't match\nwant: %s\ngot: %s", want, got)
+	}
+}