k8s/copy: check file size

I've encountered an issue in a cluster, where the tetragon bugtool files
were not complete. After debugging it seems that not all contents of the
files were copied. This patch tries to handle this problem by getting
the size of the file before the copy. If it is succesful in doing so,
then it checks whether we have read everything and if not it returns an
error.

There is already support for handling partial reads in the code, so
after this patch I am able to read the full size of the files.

Signed-off-by: Kornilios Kourtis <[email protected]>

k8s/copy.go

@@ -8,6 +8,8 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"strconv"
+	"strings"
 )
 
 const (
@@ -16,9 +18,20 @@ const (
 
 // CopyFromPod is to copy srcFile in a given pod to local destFile with defaultMaxTries.
 func (c *Client) CopyFromPod(ctx context.Context, namespace, pod, container, fromFile, destFile string, retryLimit int) error {
+
+	// attempt to get the size of the file
+	fileSize := uint64(0)
+	sizeCmd := []string{"stat", "-c", "%s", fromFile}
+	if res, err := c.ExecInPod(ctx, namespace, pod, container, sizeCmd); err == nil {
+		sizeStr := strings.TrimSuffix(res.String(), "\n")
+		if s, err := strconv.ParseUint(sizeStr, 10, 64); err == nil {
+			fileSize = s
+		}
+	}
+
 	pipe := newPipe(&CopyOptions{
 		MaxTries: retryLimit,
-		ReadFunc: readFromPod(ctx, c, namespace, pod, container, fromFile),
+		ReadFunc: readFromPod(ctx, c, namespace, pod, container, fromFile, fileSize),
 	})
 
 	outFile, err := os.OpenFile(destFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
@@ -27,20 +40,40 @@ func (c *Client) CopyFromPod(ctx context.Context, namespace, pod, container, fro
 	}
 	defer outFile.Close()
 
-	if _, err = io.Copy(outFile, pipe); err != nil {
-		return err
-	}
-	return nil
+	_, err = io.Copy(outFile, pipe)
+	return err
 }
 
-func readFromPod(ctx context.Context, client *Client, namespace, pod, container, srcFile string) ReadFunc {
+type countWriter struct {
+	total uint64
+	w     io.Writer
+}
+
+func (l *countWriter) Write(p []byte) (int, error) {
+	n, err := l.w.Write(p)
+	l.total += uint64(n)
+	return n, err
+}
+
+func readFromPod(ctx context.Context, client *Client, namespace, pod, container, srcFile string, srcFileSize uint64) ReadFunc {
+
 	return func(offset uint64, writer io.Writer) error {
 		command := []string{"sh", "-c", fmt.Sprintf(defaultReadFromByteCmd, offset, srcFile)}
-		return client.execInPodWithWriters(ctx, nil, ExecParameters{
+		outw := &countWriter{w: writer}
+		ret := client.execInPodWithWriters(ctx, nil, ExecParameters{
 			Namespace: namespace,
 			Pod:       pod,
 			Container: container,
 			Command:   command,
-		}, writer, writer)
+		}, outw, writer)
+
+		// if the read command was a success and we know the size of the source file, check
+		// whether we have read everything.
+		if ret == nil && srcFileSize > 0 {
+			if offset+outw.total < srcFileSize {
+				ret = fmt.Errorf("incomplete read: read %d from %s", offset+outw.total, srcFile)
+			}
+		}
+		return ret
 	}
 }