Make internal/ociwclayer a public API: pkg/ociwclayer, with Context for cancellation

cmd/wclayer/export.go

@@ -2,13 +2,14 @@ package main
 
 import (
 	"compress/gzip"
+	"context"
 	"io"
 	"os"
 	"path/filepath"
 
 	winio "github.com/Microsoft/go-winio"
 	"github.com/Microsoft/hcsshim/internal/appargs"
-	"github.com/Microsoft/hcsshim/internal/ociwclayer"
+	"github.com/Microsoft/hcsshim/pkg/ociwclayer"
 	"github.com/urfave/cli"
 )
 
@@ -31,13 +32,13 @@ var exportCommand = cli.Command{
 	},
 	ArgsUsage: "<layer path>",
 	Before:    appargs.Validate(appargs.NonEmptyString),
-	Action: func(context *cli.Context) (err error) {
-		path, err := filepath.Abs(context.Args().First())
+	Action: func(cliContext *cli.Context) (err error) {
+		path, err := filepath.Abs(cliContext.Args().First())
 		if err != nil {
 			return err
 		}
 
-		layers, err := normalizeLayers(context.StringSlice("layer"), true)
+		layers, err := normalizeLayers(cliContext.StringSlice("layer"), true)
 		if err != nil {
 			return err
 		}
@@ -47,7 +48,7 @@ var exportCommand = cli.Command{
 			return err
 		}
 
-		fp := context.String("output")
+		fp := cliContext.String("output")
 		f := os.Stdout
 		if fp != "" {
 			f, err = os.Create(fp)
@@ -57,10 +58,10 @@ var exportCommand = cli.Command{
 			defer f.Close()
 		}
 		w := io.Writer(f)
-		if context.Bool("gzip") {
+		if cliContext.Bool("gzip") {
 			w = gzip.NewWriter(w)
 		}
 
-		return ociwclayer.ExportLayer(w, path, layers)
+		return ociwclayer.ExportLayerToTar(context.Background(), w, path, layers)
 	},
 }

cmd/wclayer/import.go

@@ -3,13 +3,14 @@ package main
 import (
 	"bufio"
 	"compress/gzip"
+	"context"
 	"io"
 	"os"
 	"path/filepath"
 
 	"github.com/Microsoft/go-winio"
 	"github.com/Microsoft/hcsshim/internal/appargs"
-	"github.com/Microsoft/hcsshim/internal/ociwclayer"
+	"github.com/Microsoft/hcsshim/pkg/ociwclayer"
 	"github.com/urfave/cli"
 )
 
@@ -28,18 +29,18 @@ var importCommand = cli.Command{
 	},
 	ArgsUsage: "<layer path>",
 	Before:    appargs.Validate(appargs.NonEmptyString),
-	Action: func(context *cli.Context) (err error) {
-		path, err := filepath.Abs(context.Args().First())
+	Action: func(cliContext *cli.Context) (err error) {
+		path, err := filepath.Abs(cliContext.Args().First())
 		if err != nil {
 			return err
 		}
 
-		layers, err := normalizeLayers(context.StringSlice("layer"), false)
+		layers, err := normalizeLayers(cliContext.StringSlice("layer"), false)
 		if err != nil {
 			return err
 		}
 
-		fp := context.String("input")
+		fp := cliContext.String("input")
 		f := os.Stdin
 		if fp != "" {
 			f, err = os.Open(fp)
@@ -56,7 +57,7 @@ var importCommand = cli.Command{
 		if err != nil {
 			return err
 		}
-		_, err = ociwclayer.ImportLayer(r, path, layers)
+		_, err = ociwclayer.ImportLayerFromTar(context.Background(), r, path, layers)
 		return err
 	},
 }

internal/ociwclayer/export.go → pkg/ociwclayer/export.go

@@ -4,6 +4,7 @@ package ociwclayer
 
 import (
 	"archive/tar"
+	"context"
 	"io"
 	"path/filepath"
 
@@ -13,13 +14,13 @@ import (
 
 var driverInfo = hcsshim.DriverInfo{}
 
-// ExportLayer writes an OCI layer tar stream from the provided on-disk layer.
+// ExportLayerToTar writes an OCI layer tar stream from the provided on-disk layer.
 // The caller must specify the parent layers, if any, ordered from lowest to
 // highest layer.
 //
 // The layer will be mounted for this process, so the caller should ensure that
 // it is not currently mounted.
-func ExportLayer(w io.Writer, path string, parentLayerPaths []string) error {
+func ExportLayerToTar(ctx context.Context, w io.Writer, path string, parentLayerPaths []string) error {
 	err := hcsshim.ActivateLayer(driverInfo, path)
 	if err != nil {
 		return err
@@ -41,17 +42,23 @@ func ExportLayer(w io.Writer, path string, parentLayerPaths []string) error {
 		return err
 	}
 
-	err = writeTarFromLayer(r, w)
+	err = writeTarFromLayer(ctx, r, w)
 	cerr := r.Close()
 	if err != nil {
 		return err
 	}
 	return cerr
 }
 
-func writeTarFromLayer(r hcsshim.LayerReader, w io.Writer) error {
+func writeTarFromLayer(ctx context.Context, r hcsshim.LayerReader, w io.Writer) error {
 	t := tar.NewWriter(w)
 	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
+
 		name, size, fileInfo, err := r.Next()
 		if err == io.EOF {
 			break

internal/ociwclayer/import.go → pkg/ociwclayer/import.go

@@ -3,6 +3,7 @@ package ociwclayer
 import (
 	"archive/tar"
 	"bufio"
+	"context"
 	"io"
 	"os"
 	"path"
@@ -27,15 +28,15 @@ var (
 	}
 )
 
-// ImportLayer reads a layer from an OCI layer tar stream and extracts it to the
+// ImportLayerFromTar  reads a layer from an OCI layer tar stream and extracts it to the
 // specified path. The caller must specify the parent layers, if any, ordered
 // from lowest to highest layer.
 //
 // The caller must ensure that the thread or process has acquired backup and
 // restore privileges.
 //
 // This function returns the total size of the layer's files, in bytes.
-func ImportLayer(r io.Reader, path string, parentLayerPaths []string) (int64, error) {
+func ImportLayerFromTar(ctx context.Context, r io.Reader, path string, parentLayerPaths []string) (int64, error) {
 	err := os.MkdirAll(path, 0)
 	if err != nil {
 		return 0, err
@@ -44,7 +45,7 @@ func ImportLayer(r io.Reader, path string, parentLayerPaths []string) (int64, er
 	if err != nil {
 		return 0, err
 	}
-	n, err := writeLayerFromTar(r, w, path)
+	n, err := writeLayerFromTar(ctx, r, w, path)
 	cerr := w.Close()
 	if err != nil {
 		return 0, err
@@ -55,12 +56,18 @@ func ImportLayer(r io.Reader, path string, parentLayerPaths []string) (int64, er
 	return n, nil
 }
 
-func writeLayerFromTar(r io.Reader, w hcsshim.LayerWriter, root string) (int64, error) {
+func writeLayerFromTar(ctx context.Context, r io.Reader, w hcsshim.LayerWriter, root string) (int64, error) {
 	t := tar.NewReader(r)
 	hdr, err := t.Next()
 	totalSize := int64(0)
 	buf := bufio.NewWriter(nil)
 	for err == nil {
+		select {
+		case <-ctx.Done():
+			return 0, ctx.Err()
+		default:
+		}
+
 		base := path.Base(hdr.Name)
 		if strings.HasPrefix(base, whiteoutPrefix) {
 			name := path.Join(path.Dir(hdr.Name), base[len(whiteoutPrefix):])