diff --git a/images/benchmarks/profile-helper/Dockerfile b/images/benchmarks/profile-helper/Dockerfile new file mode 100644 index 000000000..e72f1f186 --- /dev/null +++ b/images/benchmarks/profile-helper/Dockerfile @@ -0,0 +1,6 @@ +FROM ubuntu:22.04 + +COPY main.go /main.go +RUN apt-get update && apt-get install -y golang-go && \ + go build -o /usr/bin/profilehelper /main.go && \ + rm -f /main.go diff --git a/images/benchmarks/profile-helper/main.go b/images/benchmarks/profile-helper/main.go new file mode 100644 index 000000000..c7e1407d7 --- /dev/null +++ b/images/benchmarks/profile-helper/main.go @@ -0,0 +1,308 @@ +// Copyright 2023 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. + +// Binary profilehelper helps exfiltrate benchmark profiles. +package main + +// This must have no non-base-library dependencies! +import ( + "archive/tar" + "compress/flate" + "crypto/sha256" + "encoding/base64" + "flag" + "fmt" + "hash" + "io" + "io/fs" + "os" + "path/filepath" + "strings" + "time" +) + +// Flags. +var ( + operation = flag.String("operation", "", "Operation to do: make-dir, delete-dir, set-containerd-flag, remove-containerd-flag, stream-dir") + flagName = flag.String("flag", "", "Name of the flag to remove or add (for containerd flag operations)") + flagValue = flag.String("value", "", "Value of the flag to add (for add-containerd-flag operations)") + containerdConfigPath = flag.String("containerd-config", "", "Path to the containerd config file (for containerd flag operations)") + containerdConfigSection = flag.String("containerd-section", "", "Name of the containerd section to modify (for containerd flag operations)") + dirPath = flag.String("dir", "", "Path to the directory to (for make-dir, delete-dir, or stream-dir operations)") + chmod = flag.Int("chmod", -1, "Permission of --dir in decimal (for chmod-dir)") +) + +// exitf aborts the program after logging the given error message to stderr. +func exitf(format string, values ...any) { + fmt.Fprintf(os.Stderr, format+"\n", values...) + os.Exit(1) +} + +// usagef aborts the program after logging the given error message and usage +// help to stderr. +func usagef(format string, values ...any) { + fmt.Fprintf(os.Stderr, format+"\n", values...) + flag.Usage() + os.Exit(2) +} + +func main() { + flag.Parse() + switch *operation { + case "": + usagef("Must specify --operation.") + case "set-containerd-flag": + setContainerdFlag( /* remove= */ false) + case "remove-containerd-flag": + setContainerdFlag( /* remove= */ true) + case "stream-dir": + streamDir() + default: + usagef("Invalid --operation: %q", *operation) + } +} + +// setContainerdFlag sets or removes a flag from the containerd config. +func setContainerdFlag(remove bool) { + if *containerdConfigPath == "" { + usagef("Must provide --containerd-config") + } + if *containerdConfigSection == "" { + usagef("Must provide --containerd-section") + } + sectionHeader := fmt.Sprintf("[%s]", *containerdConfigSection) + st, err := os.Stat(*containerdConfigPath) + if err != nil { + exitf("Cannot stat containerd config %q: %v", *containerdConfigPath, err) + } + cfg, err := os.ReadFile(*containerdConfigPath) + if err != nil { + exitf("Cannot open containerd config %q: %v", *containerdConfigPath, err) + } + isCorrectSection := false + lines := strings.Split(string(cfg), "\n") + rewritten := make([]string, 0, len(lines)) + addedFlag := false + lastFlagLine := "" + addFlag := func() { + whiteSpace := " " + if lastFlagLine != "" { + whiteSpace = lastFlagLine[:strings.Index(lastFlagLine, strings.TrimSpace(lastFlagLine))] + } + rewritten = append(rewritten, fmt.Sprintf("%s%s = %q", whiteSpace, *flagName, *flagValue)) + } + for _, line := range lines { + if line == "" { + continue + } + if line == sectionHeader { + isCorrectSection = true + rewritten = append(rewritten, line) + continue + } + if !isCorrectSection { + rewritten = append(rewritten, line) + continue + } + if isCorrectSection && strings.HasPrefix(line, "[") { + if !remove && !addedFlag { + // Need to add the flag before we close out this section. + addFlag() + } + isCorrectSection = false + rewritten = append(rewritten, line) + break + } + equalIndex := strings.IndexRune(line, '=') + if equalIndex == -1 { + rewritten = append(rewritten, line) + continue + } + lastFlagLine = line + if f := strings.TrimSpace(line[:equalIndex]); f != *flagName { + rewritten = append(rewritten, line) + continue + } + if remove || addedFlag { + // Omit this line entirely. + continue + } + // Replace this line. + addedFlag = true + // Use `line[:equalIndex]` to preserve correct indentation. + rewritten = append(rewritten, fmt.Sprintf("%s= %q", line[:equalIndex], *flagValue)) + } + if !remove && !addedFlag { + addFlag() + } + + newCfg := strings.TrimSpace(strings.Join(rewritten, "\n")) + "\n" + if strings.TrimSpace(newCfg) == strings.TrimSpace(string(cfg)) { + fmt.Fprintf(os.Stderr, "Containerd config is unchanged:\n%s\n\n", string(cfg)) + } + fmt.Fprintf(os.Stderr, "Previous containerd config:\n%s\n\nNew containerd config:\n%s\n\n", string(cfg), newCfg) + if err := os.WriteFile(*containerdConfigPath, []byte(newCfg), st.Mode()); err != nil { + exitf("Cannot write to containerd config %q: %v", *containerdConfigPath, err) + } +} + +// dirStreamWriter handles stream-dir write requests. +// It is the last writer in the chain, after base64 encoding. +// Its role is to split up the base64 stream into lines, prefix it +// with "DATA:", and to rate-limit it. It also keeps track of the +// hash of the stream. +type dirStreamWriter struct { + beganStream bool + dirPath string + lastLine time.Time + checksum hash.Hash +} + +// Write implements `io.WriteCloser.Write`. +func (w *dirStreamWriter) Write(b []byte) (int, error) { + const ( + maxLineSize = 384 + linesPerSecond = 8 + durationPerLine = time.Second / time.Duration(linesPerSecond) + ) + if !w.beganStream { + if _, err := fmt.Fprintf(os.Stdout, "BEGIN:%s\n", w.dirPath); err != nil { + return 0, err + } + w.beganStream = true + } + var written int + for i := 0; i < len(b); i += maxLineSize { + if sinceLastLine := time.Since(w.lastLine); sinceLastLine < durationPerLine { + time.Sleep(durationPerLine - sinceLastLine) + } + w.lastLine = time.Now() + _, err := os.Stdout.WriteString("DATA:") + // We don't increase `written` here, because the write function's + // bytes written should represent the bytes consumed from `b`. + if err != nil { + return written, err + } + endIndex := i + maxLineSize + if endIndex > len(b) { + endIndex = len(b) + } + n, err := os.Stdout.Write(b[i:endIndex]) + if n > 0 { + written += n + w.checksum.Write(b[i : i+n]) + } + if err != nil { + return written, err + } + _, err = os.Stdout.WriteString("\n") + // We don't increase `written` here for the same reason. + if err != nil { + return written, err + } + } + return written, nil +} + +// Close implements `io.WriteCloser.Close`. +func (w *dirStreamWriter) Close() error { + if !w.beganStream { + return nil + } + _, err := fmt.Fprintf(os.Stdout, "SHA256:%x\n", w.checksum.Sum(nil)) + if err != nil { + return fmt.Errorf("cannot write SHA256: %w", err) + } + return nil +} + +// streamDir tars up all the files in a target directory, then +// compresses it with flate, then base64 the compressed data, and then +// outputs it to stdout across multiple lines. +// Its goal is not to be efficient, but rather to work in very constrained +// environments where the only effective data exfiltration mechanism is +// a single standard buffered stdout log stream using printable ASCII +// characters only. +// Each line of stdout is prefixed with either "BEGIN:", "DATA:", +// or "SHA256:" (marking the end of the stream). +// The program will wait for 15 seconds before streaming any output; +// this gives time for a client to attach to the log stream before anything +// is output. This needs to happen before any output occurs, in case the +// log buffer isn't large enough to contain the entire directory's contents. +func streamDir() { + if *dirPath == "" { + usagef("Must provide --dir") + } + dsw := &dirStreamWriter{ + dirPath: *dirPath, + checksum: sha256.New(), + } + enc := base64.NewEncoder(base64.StdEncoding, dsw) + fl, err := flate.NewWriter(enc, flate.BestCompression) + if err != nil { + exitf("Cannot create flate writer: %v", err) + } + tw := tar.NewWriter(fl) + time.Sleep(15 * time.Second) + err = filepath.Walk(*dirPath, func(path string, info fs.FileInfo, walkErr error) error { + if walkErr != nil { + // Keep walking other directories, so return nil here. + return nil + } + if info.IsDir() { + return nil + } + size := info.Size() + clean, err := filepath.Rel(*dirPath, path) + if err != nil { + return fmt.Errorf("cannot express path %q relative to base path %q: %w", path, *dirPath, err) + } + f, err := os.Open(path) + if err != nil { + return fmt.Errorf("cannot open %q: %w", path, err) + } + hdr := tar.Header{ + Name: clean, + Mode: int64(info.Mode()), + Size: size, + Format: tar.FormatPAX, + } + if err := tw.WriteHeader(&hdr); err != nil { + return fmt.Errorf("cannot write tar header for %q: %w", path, err) + } + copied, err := io.Copy(tw, f) + if err != nil { + return fmt.Errorf("cannot write tar body for %q: %w", path, err) + } + if copied != size { + return fmt.Errorf("file size for %q was inaccurate: stat claimed %d bytes, copied %d", path, size, copied) + } + return nil + }) + if err != nil { + exitf("Cannot walk directory %q: %v", *dirPath, err) + } + if err := tw.Close(); err != nil { + exitf("Cannot close tar file: %v", err) + } + if err := fl.Close(); err != nil { + exitf("Cannot close flate writer: %v", err) + } + if err := enc.Close(); err != nil { + exitf("Cannot close base64 encoder: %v", err) + } + if err := dsw.Close(); err != nil { + exitf("Cannot close directory stream: %v", err) + } +}