Merge pull request #9146 from sitano:ivan_nocompressio

PiperOrigin-RevId: 555750100

pkg/compressio/BUILD

@@ -7,7 +7,10 @@ package(
 
 go_library(
     name = "compressio",
-    srcs = ["compressio.go"],
+    srcs = [
+        "compressio.go",
+        "nocompressio.go",
+    ],
     visibility = ["//:sandbox"],
     deps = ["//pkg/sync"],
 )
@@ -15,6 +18,9 @@ go_library(
 go_test(
     name = "compressio_test",
     size = "medium",
-    srcs = ["compressio_test.go"],
+    srcs = [
+        "compressio_test.go",
+        "nocompressio_test.go",
+    ],
     library = ":compressio",
 )

pkg/compressio/compressio.go

@@ -13,7 +13,8 @@
 // limitations under the License.
 
 // Package compressio provides parallel compression and decompression, as well
-// as optional SHA-256 hashing.
+// as optional SHA-256 hashing. It also provides another storage variant
+// (nocompressio) that does not compress data but tracks its integrity.
 //
 // The stream format is defined as follows.
 //

pkg/compressio/compressio_test.go

@@ -96,6 +96,11 @@ func doTest(t harness, opts testOpts) {
 	compressionTime := time.Since(compressionStartTime)
 	compressionRatio := float32(compressed.Len()) / float32(len(opts.Data))
 
+	if compressed.Len() == 0 {
+		// Data can't be corrupted if there is no data.
+		opts.CorruptData = false
+	}
+
 	// Decompress.
 	var decompressed bytes.Buffer
 	decompressionStartTime := time.Now()

pkg/compressio/nocompressio.go

@@ -0,0 +1,283 @@
+// 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.
+
+package compressio
+
+import (
+	"bufio"
+	"bytes"
+	"crypto/hmac"
+	"crypto/sha256"
+	"encoding/binary"
+	"hash"
+	"io"
+)
+
+// nocompressio provides data storage that does not use data compression but
+// offers optional data integrity via SHA-256 hashing.
+//
+// The stream format is defined as follows.
+//
+// /------------------------------------------------------\
+// |                  data size (4-bytes)                 |
+// +------------------------------------------------------+
+// |                  data                                |
+// +------------------------------------------------------+
+// |       (optional) hash (32-bytes)                     |
+// +------------------------------------------------------+
+// |                  data size (4-bytes)                 |
+// +------------------------------------------------------+
+// |                       ......                         |
+// \------------------------------------------------------/
+//
+// where each hash is calculated from the following items in order
+//
+//	data
+//	data size
+
+// SimpleReader is a reader from uncompressed image.
+type SimpleReader struct {
+	// in is the source.
+	in io.Reader
+
+	// key is the key used to create hash objects.
+	key []byte
+
+	// h is the hash object.
+	h hash.Hash
+
+	// current data chunk size
+	chunkSize uint32
+
+	// current chunk position
+	done uint32
+}
+
+var _ io.Reader = (*SimpleReader)(nil)
+
+const (
+	defaultBufSize = 256 * 1024
+)
+
+// NewSimpleReader returns a new (uncompressed) reader. If key is non-nil, the data stream
+// is assumed to contain expected hash values. See package comments for
+// details.
+func NewSimpleReader(in io.Reader, key []byte) (*SimpleReader, error) {
+	r := &SimpleReader{
+		in:  bufio.NewReaderSize(in, defaultBufSize),
+		key: key,
+	}
+
+	if key != nil {
+		r.h = hmac.New(sha256.New, key)
+	}
+
+	return r, nil
+}
+
+// ReadByte implements wire.Reader.ReadByte.
+func (r *SimpleReader) ReadByte() (byte, error) {
+	var p [1]byte
+	n, err := r.Read(p[:])
+	if n != 1 {
+		return p[0], err
+	}
+	// Suppress EOF.
+	return p[0], nil
+}
+
+// Read implements io.Reader.Read.
+func (r *SimpleReader) Read(p []byte) (int, error) {
+	var scratch [4]byte
+
+	if len(p) == 0 {
+		return r.in.Read(p)
+	}
+
+	// need next chunk?
+	if r.done >= r.chunkSize {
+		if _, err := io.ReadFull(r.in, scratch[:]); err != nil {
+			return 0, err
+		}
+
+		r.chunkSize = binary.BigEndian.Uint32(scratch[:])
+		r.done = 0
+		if r.key != nil {
+			r.h.Reset()
+		}
+
+		if r.chunkSize == 0 {
+			// this must not happen
+			return 0, io.ErrNoProgress
+		}
+	}
+
+	toRead := uint32(len(p))
+	// can't read more than whats left
+	if toRead > r.chunkSize-r.done {
+		toRead = r.chunkSize - r.done
+	}
+
+	n, err := r.in.Read(p[:toRead])
+	if err != nil {
+		if err == io.EOF {
+			// this only can happen if storage or data size is corrupted,
+			// but we have no other means to detect it earlier as we store
+			// hash after the data block.
+			return n, ErrHashMismatch
+		}
+		return n, err
+	}
+
+	if r.key != nil {
+		_, _ = r.h.Write(p[:n])
+	}
+
+	r.done += uint32(n)
+	if r.done >= r.chunkSize {
+		if r.key != nil {
+			binary.BigEndian.PutUint32(scratch[:], r.chunkSize)
+			r.h.Write(scratch[:4])
+
+			sum := r.h.Sum(nil)
+			readerSum := make([]byte, len(sum))
+			if _, err := io.ReadFull(r.in, readerSum); err != nil {
+				if err == io.EOF {
+					return n, io.ErrUnexpectedEOF
+				}
+				return n, err
+			}
+
+			if !hmac.Equal(readerSum, sum) {
+				return n, ErrHashMismatch
+			}
+		}
+
+		r.done = 0
+		r.chunkSize = 0
+	}
+
+	return n, nil
+}
+
+// SimpleWriter is a writer that does not compress.
+type SimpleWriter struct {
+	// base is the underlying writer.
+	base io.Writer
+
+	// out is a buffered writer.
+	out *bufio.Writer
+
+	// key is the key used to create hash objects.
+	key []byte
+
+	// closed indicates whether the file has been closed.
+	closed bool
+}
+
+var _ io.Writer = (*SimpleWriter)(nil)
+var _ io.Closer = (*SimpleWriter)(nil)
+
+// NewSimpleWriter returns a new non-compressing writer. If key is non-nil, hash values are
+// generated and written out for compressed bytes. See package comments for
+// details.
+func NewSimpleWriter(out io.Writer, key []byte) (*SimpleWriter, error) {
+	return &SimpleWriter{
+		base: out,
+		out:  bufio.NewWriterSize(out, defaultBufSize),
+		key:  key,
+	}, nil
+}
+
+// WriteByte implements wire.Writer.WriteByte.
+//
+// Note that this implementation is necessary on the object itself, as an
+// interface-based dispatch cannot tell whether the array backing the slice
+// escapes, therefore the all bytes written will generate an escape.
+func (w *SimpleWriter) WriteByte(b byte) error {
+	var p [1]byte
+	p[0] = b
+	n, err := w.Write(p[:])
+	if n != 1 {
+		return err
+	}
+	return nil
+}
+
+// Write implements io.Writer.Write.
+func (w *SimpleWriter) Write(p []byte) (int, error) {
+	var scratch [4]byte
+
+	// Did we close already?
+	if w.closed {
+		return 0, io.ErrUnexpectedEOF
+	}
+
+	l := uint32(len(p))
+
+	// chunk length
+	binary.BigEndian.PutUint32(scratch[:], l)
+	if _, err := w.out.Write(scratch[:4]); err != nil {
+		return 0, err
+	}
+
+	// Write out to the stream.
+	n, err := w.out.Write(p)
+	if err != nil {
+		return n, err
+	}
+
+	if w.key != nil {
+		h := hmac.New(sha256.New, w.key)
+
+		// chunk data
+		_, _ = h.Write(p)
+
+		// chunk length
+		binary.BigEndian.PutUint32(scratch[:], l)
+		h.Write(scratch[:4])
+
+		sum := h.Sum(nil)
+		if _, err := io.CopyN(w.out, bytes.NewReader(sum), int64(len(sum))); err != nil {
+			return n, err
+		}
+	}
+
+	return n, nil
+}
+
+// Close implements io.Closer.Close.
+func (w *SimpleWriter) Close() error {
+	// Did we already close? After the call to Close, we always mark as
+	// closed, regardless of whether the flush is successful.
+	if w.closed {
+		return io.ErrUnexpectedEOF
+	}
+	w.closed = true
+
+	// Flush buffered writer
+	if err := w.out.Flush(); err != nil {
+		return err
+	}
+
+	// Close the underlying writer (if necessary).
+	if closer, ok := w.base.(io.Closer); ok {
+		return closer.Close()
+	}
+
+	w.out = nil
+	w.base = nil
+
+	return nil
+}

pkg/compressio/nocompressio_test.go

@@ -0,0 +1,68 @@
+// 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.
+
+package compressio
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"math/rand"
+	"testing"
+	"time"
+)
+
+func TestNoCompress(t *testing.T) {
+	rand.Seed(time.Now().Unix())
+
+	var (
+		data  = initTest(t, 10*1024*1024)
+		data0 = data[:0]
+		data1 = data[:1]
+		data2 = data[:11]
+		data3 = data[:16]
+		data4 = data[:]
+	)
+
+	for _, data := range [][]byte{data0, data1, data2, data3, data4} {
+		for _, blockSize := range []uint32{1, 4, 1024, 4 * 1024, 16 * 1024} {
+			// Skip annoying tests; they just take too long.
+			if blockSize <= 16 && len(data) > 16 {
+				continue
+			}
+
+			for _, key := range [][]byte{nil, hashKey} {
+				for _, corruptData := range []bool{false, true} {
+					if key == nil && corruptData {
+						// No need to test corrupt data
+						// case when not doing hashing.
+						continue
+					}
+					// Do the compress test.
+					doTest(t, testOpts{
+						Name: fmt.Sprintf("len(data)=%d, blockSize=%d, key=%s, corruptData=%v", len(data), blockSize, string(key), corruptData),
+						Data: data,
+						NewWriter: func(b *bytes.Buffer) (io.Writer, error) {
+							return NewSimpleWriter(b, key)
+						},
+						NewReader: func(b *bytes.Buffer) (io.Reader, error) {
+							return NewSimpleReader(b, key)
+						},
+						CorruptData: corruptData,
+					})
+				}
+			}
+		}
+	}
+}

pkg/state/statefile/statefile.go

@@ -84,6 +84,68 @@ var ErrInvalidMetadataLength = fmt.Errorf("metadata length invalid, maximum size
 // ErrMetadataInvalid is returned if passed metadata is invalid.
 var ErrMetadataInvalid = fmt.Errorf("metadata invalid, can't start with _")
 
+// ErrInvalidFlags is returned if passed flags set is invalid.
+var ErrInvalidFlags = fmt.Errorf("flags set is invalid")
+
+const (
+	compressionKey = "compression"
+)
+
+// CompressionLevel is the image compression level.
+type CompressionLevel string
+
+const (
+	// CompressionLevelFlateBestSpeed represents flate algorithm in best-speed mode.
+	CompressionLevelFlateBestSpeed = CompressionLevel("flate-best-speed")
+	// CompressionLevelNone represents the absence of any compression on an image.
+	CompressionLevelNone = CompressionLevel("none")
+)
+
+// Options is statefile options.
+type Options struct {
+	// Compression is an image compression type/level.
+	Compression CompressionLevel
+}
+
+// WriteToMetadata save options to the metadata storage.  Method returns the
+// reference to the original metadata map to allow to be used in the chain calls.
+func (o Options) WriteToMetadata(metadata map[string]string) map[string]string {
+	metadata[compressionKey] = string(o.Compression)
+	return metadata
+}
+
+// CompressionLevelFromString parses a string into the CompressionLevel.
+func CompressionLevelFromString(val string) (CompressionLevel, error) {
+	switch val {
+	case string(CompressionLevelFlateBestSpeed):
+		return CompressionLevelFlateBestSpeed, nil
+	case string(CompressionLevelNone):
+		return CompressionLevelNone, nil
+	default:
+		return CompressionLevelNone, ErrInvalidFlags
+	}
+}
+
+// CompressionLevelFromMetadata returns image compression type stored in the metadata.
+// If the metadata doesn't contain compression information the default behavior
+// is the "flate-best-speed" state because the default behavior used to be to always
+// compress. If the parameter is missing it will be set to default.
+func CompressionLevelFromMetadata(metadata map[string]string) (CompressionLevel, error) {
+	var err error
+
+	compression := CompressionLevelFlateBestSpeed
+
+	if val, ok := metadata[compressionKey]; ok {
+		if compression, err = CompressionLevelFromString(val); err != nil {
+			return CompressionLevelNone, err
+		}
+	} else {
+		metadata[compressionKey] = string(compression)
+	}
+
+	return compression, nil
+}
+
 // WriteCloser is an io.Closer and wire.Writer.
 type WriteCloser interface {
 	wire.Writer
@@ -123,6 +185,12 @@ func NewWriter(w io.Writer, key []byte, metadata map[string]string) (WriteCloser
 	metadata["_timestamp"] = time.Now().UTC().String()
 	defer delete(metadata, "_timestamp")
 
+	// Save compression state
+	compression, err := CompressionLevelFromMetadata(metadata)
+	if err != nil {
+		return nil, err
+	}
+
 	// Write the metadata.
 	b, err := json.Marshal(metadata)
 	if err != nil {
@@ -152,11 +220,15 @@ func NewWriter(w io.Writer, key []byte, metadata map[string]string) (WriteCloser
 		}
 	}
 
-	// Wrap in compression. We always use "best speed" mode here. When using
-	// "best compression" mode, there is usually only a little gain in file
-	// size reduction, which translate to even smaller gain in restore
-	// latency reduction, while inccuring much more CPU usage at save time.
-	return compressio.NewWriter(w, key, compressionChunkSize, flate.BestSpeed)
+	// Wrap in compression. When using "best compression" mode, there is usually
+	// only a little gain in file size reduction, which translate to even smaller
+	// gain in restore latency reduction, while inccuring much more CPU usage at
+	// save time.
+	if compression == CompressionLevelFlateBestSpeed {
+		return compressio.NewWriter(w, key, compressionChunkSize, flate.BestSpeed)
+	}
+
+	return compressio.NewSimpleWriter(w, key)
 }
 
 // MetadataUnsafe reads out the metadata from a state file without verifying any
@@ -245,10 +317,29 @@ func NewReader(r io.Reader, key []byte) (wire.Reader, map[string]string, error)
 		return nil, nil, err
 	}
 
-	// Wrap in compression.
-	cr, err := compressio.NewReader(r, key)
+	// Determine image compression state. If the metadata doesn't contain
+	// compression information the default behavior is the "compressed" state
+	// because the default behavior used to be to always compress.
+	compression, err := CompressionLevelFromMetadata(metadata)
 	if err != nil {
 		return nil, nil, err
 	}
+
+	// Pick correct reader
+	var cr wire.Reader
+
+	if compression == CompressionLevelFlateBestSpeed {
+		cr, err = compressio.NewReader(r, key)
+	} else if compression == CompressionLevelNone {
+		cr, err = compressio.NewSimpleReader(r, key)
+	} else {
+		// Should never occur, as it has the default path.
+		return nil, nil, fmt.Errorf("metadata contains invalid compression flag value: %v", compression)
+	}
+
+	if err != nil {
+		return nil, nil, err
+	}
+
 	return cr, metadata, nil
 }

pkg/state/statefile/statefile_test.go

@@ -46,6 +46,11 @@ type testCase struct {
 func TestStatefile(t *testing.T) {
 	rand.Seed(time.Now().Unix())
 
+	compression := map[string]CompressionLevel{
+		"none":       CompressionLevelNone,
+		"compressed": CompressionLevelFlateBestSpeed,
+	}
+
 	cases := []testCase{
 		// Various data sizes.
 		{"nil", nil, nil},
@@ -72,90 +77,102 @@ func TestStatefile(t *testing.T) {
 		{"two metadata", []byte("data"), map[string]string{"foo": "bar", "one": "two"}},
 	}
 
-	for _, c := range cases {
-		// Generate a key.
-		integrityKey, err := randomKey()
-		if err != nil {
-			t.Errorf("can't generate key: got %v, excepted nil", err)
-			continue
-		}
+	for cKey, compress := range compression {
+		t.Run(cKey, func(t *testing.T) {
+			for _, c := range cases {
+				// Generate a key.
+				integrityKey, err := randomKey()
+				if err != nil {
+					t.Errorf("can't generate key: got %v, excepted nil", err)
+					continue
+				}
 
-		t.Run(c.name, func(t *testing.T) {
-			for _, key := range [][]byte{nil, integrityKey} {
-				t.Run("key="+string(key), func(t *testing.T) {
-					// Encoding happens via a buffer.
-					var bufEncoded bytes.Buffer
-					var bufDecoded bytes.Buffer
+				// Save compression state
+				if c.metadata == nil {
+					c.metadata = map[string]string{}
+				}
 
-					// Do all the writing.
-					w, err := NewWriter(&bufEncoded, key, c.metadata)
-					if err != nil {
-						t.Fatalf("error creating writer: got %v, expected nil", err)
-					}
-					if _, err := io.Copy(w, bytes.NewBuffer(c.data)); err != nil {
-						t.Fatalf("error during write: got %v, expected nil", err)
-					}
+				c.metadata[compressionKey] = string(compress)
 
-					// Finish the sum.
-					if err := w.Close(); err != nil {
-						t.Fatalf("error during close: got %v, expected nil", err)
-					}
+				t.Run(c.name, func(t *testing.T) {
+					for _, key := range [][]byte{nil, integrityKey} {
+						t.Run("key="+string(key), func(t *testing.T) {
+							// Encoding happens via a buffer.
+							var bufEncoded bytes.Buffer
+							var bufDecoded bytes.Buffer
 
-					t.Logf("original data: %d bytes, encoded: %d bytes.",
-						len(c.data), len(bufEncoded.Bytes()))
+							// Do all the writing.
+							w, err := NewWriter(&bufEncoded, key, c.metadata)
+							if err != nil {
+								t.Fatalf("error creating writer: got %v, expected nil", err)
+							}
+							if _, err := io.Copy(w, bytes.NewBuffer(c.data)); err != nil {
+								t.Fatalf("error during write: got %v, expected nil", err)
+							}
 
-					// Do all the reading.
-					r, metadata, err := NewReader(bytes.NewReader(bufEncoded.Bytes()), key)
-					if err != nil {
-						t.Fatalf("error creating reader: got %v, expected nil", err)
-					}
-					if _, err := io.Copy(&bufDecoded, r); err != nil {
-						t.Fatalf("error during read: got %v, expected nil", err)
-					}
+							// Finish the sum.
+							if err := w.Close(); err != nil {
+								t.Fatalf("error during close: got %v, expected nil", err)
+							}
 
-					// Check that the data matches.
-					if !bytes.Equal(c.data, bufDecoded.Bytes()) {
-						t.Fatalf("data didn't match (%d vs %d bytes)", len(bufDecoded.Bytes()), len(c.data))
-					}
+							t.Logf("original data: %d bytes, encoded: %d bytes.",
+								len(c.data), len(bufEncoded.Bytes()))
 
-					// Check that the metadata matches.
-					for k, v := range c.metadata {
-						nv, ok := metadata[k]
-						if !ok {
-							t.Fatalf("missing metadata: %s", k)
-						}
-						if v != nv {
-							t.Fatalf("mismatched metdata for %s: got %s, expected %s", k, nv, v)
-						}
-					}
+							// Do all the reading.
+							r, metadata, err := NewReader(bytes.NewReader(bufEncoded.Bytes()), key)
+							if err != nil {
+								t.Fatalf("error creating reader: got %v, expected nil", err)
+							}
+							if _, err := io.Copy(&bufDecoded, r); err != nil {
+								t.Fatalf("error during read: got %v, expected nil", err)
+							}
 
-					// Change the data and verify that it fails.
-					if key != nil {
-						b := append([]byte(nil), bufEncoded.Bytes()...)
-						b[rand.Intn(len(b))]++
-						bufDecoded.Reset()
-						r, _, err = NewReader(bytes.NewReader(b), key)
-						if err == nil {
-							_, err = io.Copy(&bufDecoded, r)
-						}
-						if err == nil {
-							t.Error("got no error: expected error on data corruption")
-						}
-					}
+							// Check that the data matches.
+							if !bytes.Equal(c.data, bufDecoded.Bytes()) {
+								t.Fatalf("data didn't match (%d vs %d bytes)", len(bufDecoded.Bytes()), len(c.data))
+							}
 
-					// Change the key and verify that it fails.
-					newKey := integrityKey
-					if len(key) > 0 {
-						newKey = append([]byte{}, key...)
-						newKey[rand.Intn(len(newKey))]++
-					}
-					bufDecoded.Reset()
-					r, _, err = NewReader(bytes.NewReader(bufEncoded.Bytes()), newKey)
-					if err == nil {
-						_, err = io.Copy(&bufDecoded, r)
-					}
-					if err != compressio.ErrHashMismatch {
-						t.Errorf("got error: %v, expected ErrHashMismatch on key mismatch", err)
+							// Check that the metadata matches.
+							for k, v := range c.metadata {
+								nv, ok := metadata[k]
+								if !ok {
+									t.Fatalf("missing metadata: %s", k)
+								}
+								if v != nv {
+									t.Fatalf("mismatched metdata for %s: got %s, expected %s", k, nv, v)
+								}
+							}
+
+							// Change the data and verify that it fails.
+							if key != nil {
+								b := append([]byte(nil), bufEncoded.Bytes()...)
+								i := rand.Intn(len(b))
+								b[i]++
+								bufDecoded.Reset()
+								r, _, err = NewReader(bytes.NewReader(b), key)
+								if err == nil {
+									_, err = io.Copy(&bufDecoded, r)
+								}
+								if err == nil {
+									t.Errorf("got no error: expected error on data corruption in byte [%d] = %x", i, b[i])
+								}
+							}
+
+							// Change the key and verify that it fails.
+							newKey := integrityKey
+							if len(key) > 0 {
+								newKey = append([]byte{}, key...)
+								newKey[rand.Intn(len(newKey))]++
+							}
+							bufDecoded.Reset()
+							r, _, err = NewReader(bytes.NewReader(bufEncoded.Bytes()), newKey)
+							if err == nil {
+								_, err = io.Copy(&bufDecoded, r)
+							}
+							if err != compressio.ErrHashMismatch {
+								t.Errorf("got error: %v, expected ErrHashMismatch on key mismatch", err)
+							}
+						})
 					}
 				})
 			}
@@ -198,7 +215,7 @@ func benchmark(b *testing.B, size int, write bool, compressible bool) {
 	var stateBuf bytes.Buffer
 	writeState := func() {
 		stateBuf.Reset()
-		w, err := NewWriter(&stateBuf, key, nil)
+		w, err := NewWriter(&stateBuf, key, Options{Compression: CompressionLevelFlateBestSpeed}.WriteToMetadata(map[string]string{}))
 		if err != nil {
 			b.Fatalf("error creating writer: %v", err)
 		}

runsc/cmd/checkpoint.go

@@ -16,12 +16,14 @@ package cmd
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"path/filepath"
 
 	"github.com/google/subcommands"
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/log"
+	"gvisor.dev/gvisor/pkg/state/statefile"
 	"gvisor.dev/gvisor/runsc/cmd/util"
 	"gvisor.dev/gvisor/runsc/config"
 	"gvisor.dev/gvisor/runsc/container"
@@ -36,6 +38,7 @@ const checkpointFileName = "checkpoint.img"
 type Checkpoint struct {
 	imagePath    string
 	leaveRunning bool
+	compression  CheckpointCompression
 }
 
 // Name implements subcommands.Command.Name.
@@ -58,6 +61,7 @@ func (*Checkpoint) Usage() string {
 func (c *Checkpoint) SetFlags(f *flag.FlagSet) {
 	f.StringVar(&c.imagePath, "image-path", "", "directory path to saved container image")
 	f.BoolVar(&c.leaveRunning, "leave-running", false, "restart the container after checkpointing")
+	f.Var(newCheckpointCompressionValue(statefile.CompressionLevelFlateBestSpeed, &c.compression), "compression", "compress checkpoint image on disk. Values: none|flate-best-speed.")
 
 	// Unimplemented flags necessary for compatibility with docker.
 	var wp string
@@ -97,7 +101,7 @@ func (c *Checkpoint) Execute(_ context.Context, f *flag.FlagSet, args ...any) su
 	}
 	defer file.Close()
 
-	if err := cont.Checkpoint(file); err != nil {
+	if err := cont.Checkpoint(file, statefile.Options{Compression: c.compression.Level()}); err != nil {
 		util.Fatalf("checkpoint failed: %v", err)
 	}
 
@@ -156,3 +160,40 @@ func (c *Checkpoint) Execute(_ context.Context, f *flag.FlagSet, args ...any) su
 
 	return subcommands.ExitSuccess
 }
+
+// CheckpointCompression represents checkpoint image writer behavior. The
+// default behavior is to compress because the default behavior used to be to
+// always compress.
+type CheckpointCompression statefile.CompressionLevel
+
+func newCheckpointCompressionValue(val statefile.CompressionLevel, p *CheckpointCompression) *CheckpointCompression {
+	*p = CheckpointCompression(val)
+	return (*CheckpointCompression)(p)
+}
+
+// Set implements flag.Value.
+func (g *CheckpointCompression) Set(v string) error {
+	t, err := statefile.CompressionLevelFromString(v)
+	if err != nil {
+		return fmt.Errorf("invalid checkpoint compression type %q", v)
+	}
+
+	*g = CheckpointCompression(t)
+
+	return nil
+}
+
+// Get implements flag.Getter.
+func (g *CheckpointCompression) Get() any {
+	return *g
+}
+
+// String implements flag.Value.
+func (g CheckpointCompression) String() string {
+	return string(g)
+}
+
+// Level returns corresponding statefile.CompressionLevel value.
+func (g CheckpointCompression) Level() statefile.CompressionLevel {
+	return statefile.CompressionLevel(g)
+}

runsc/container/BUILD

@@ -24,6 +24,7 @@ go_library(
         "//pkg/sentry/control",
         "//pkg/sentry/pgalloc",
         "//pkg/sighandling",
+        "//pkg/state/statefile",
         "//pkg/sync",
         "//runsc/boot",
         "//runsc/cgroup",
@@ -80,6 +81,7 @@ go_test(
         "//pkg/sentry/seccheck",
         "//pkg/sentry/seccheck/points:points_go_proto",
         "//pkg/sentry/seccheck/sinks/remote/test",
+        "//pkg/state/statefile",
         "//pkg/sync",
         "//pkg/test/testutil",
         "//pkg/unet",

runsc/container/container.go

@@ -39,6 +39,7 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/control"
 	"gvisor.dev/gvisor/pkg/sentry/pgalloc"
 	"gvisor.dev/gvisor/pkg/sighandling"
+	"gvisor.dev/gvisor/pkg/state/statefile"
 	"gvisor.dev/gvisor/runsc/boot"
 	"gvisor.dev/gvisor/runsc/cgroup"
 	"gvisor.dev/gvisor/runsc/config"
@@ -691,12 +692,12 @@ func (c *Container) ForwardSignals(pid int32, fgProcess bool) func() {
 
 // Checkpoint sends the checkpoint call to the container.
 // The statefile will be written to f, the file at the specified image-path.
-func (c *Container) Checkpoint(f *os.File) error {
+func (c *Container) Checkpoint(f *os.File, options statefile.Options) error {
 	log.Debugf("Checkpoint container, cid: %s", c.ID)
 	if err := c.requireStatus("checkpoint", Created, Running, Paused); err != nil {
 		return err
 	}
-	return c.Sandbox.Checkpoint(c.ID, f)
+	return c.Sandbox.Checkpoint(c.ID, f, options)
 }
 
 // Pause suspends the container and its kernel.

runsc/container/container_test.go

@@ -40,6 +40,7 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/kernel"
 	"gvisor.dev/gvisor/pkg/sentry/kernel/auth"
 	"gvisor.dev/gvisor/pkg/sentry/platform"
+	"gvisor.dev/gvisor/pkg/state/statefile"
 	"gvisor.dev/gvisor/pkg/sync"
 	"gvisor.dev/gvisor/pkg/test/testutil"
 	"gvisor.dev/gvisor/runsc/boot"
@@ -1076,7 +1077,7 @@ func TestCheckpointRestore(t *testing.T) {
 			}
 
 			// Checkpoint running container; save state into new file.
-			if err := cont.Checkpoint(file); err != nil {
+			if err := cont.Checkpoint(file, statefile.Options{Compression: statefile.CompressionLevelFlateBestSpeed}); err != nil {
 				t.Fatalf("error checkpointing container to empty file: %v", err)
 			}
 			defer os.RemoveAll(imagePath)
@@ -1255,7 +1256,7 @@ func TestUnixDomainSockets(t *testing.T) {
 			}
 
 			// Checkpoint running container; save state into new file.
-			if err := cont.Checkpoint(file); err != nil {
+			if err := cont.Checkpoint(file, statefile.Options{Compression: statefile.CompressionLevelFlateBestSpeed}); err != nil {
 				t.Fatalf("error checkpointing container to empty file: %v", err)
 			}
 

runsc/sandbox/BUILD

@@ -32,6 +32,7 @@ go_library(
         "//pkg/sentry/control",
         "//pkg/sentry/platform",
         "//pkg/sentry/seccheck",
+        "//pkg/state/statefile",
         "//pkg/sync",
         "//pkg/tcpip/header",
         "//pkg/tcpip/stack",

runsc/sandbox/sandbox.go

@@ -45,6 +45,7 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/control"
 	"gvisor.dev/gvisor/pkg/sentry/platform"
 	"gvisor.dev/gvisor/pkg/sentry/seccheck"
+	"gvisor.dev/gvisor/pkg/state/statefile"
 	"gvisor.dev/gvisor/pkg/sync"
 	"gvisor.dev/gvisor/pkg/urpc"
 	"gvisor.dev/gvisor/runsc/boot"
@@ -1217,9 +1218,10 @@ func (s *Sandbox) SignalProcess(cid string, pid int32, sig unix.Signal, fgProces
 
 // Checkpoint sends the checkpoint call for a container in the sandbox.
 // The statefile will be written to f.
-func (s *Sandbox) Checkpoint(cid string, f *os.File) error {
-	log.Debugf("Checkpoint sandbox %q", s.ID)
+func (s *Sandbox) Checkpoint(cid string, f *os.File, options statefile.Options) error {
+	log.Debugf("Checkpoint sandbox %q, options %+v", s.ID, options)
 	opt := control.SaveOpts{
+		Metadata: options.WriteToMetadata(map[string]string{}),
 		FilePayload: urpc.FilePayload{
 			Files: []*os.File{f},
 		},