github.com/microsoft/moc@v0.17.1/pkg/logging/filetransfer.go (about)

     1  // Copyright (c) Microsoft Corporation.
     2  // Licensed under the Apache v2.0 license.
     3  
     4  package logging
     5  
     6  import (
     7  	"context"
     8  	"io"
     9  	"os"
    10  )
    11  
    12  const BUFFER_SIZE = 1024
    13  
    14  func UploadFile(ctx context.Context, filename string, sendFunc func([]byte, error) error) error {
    15  	f, err := os.Open(filename)
    16  	if err != nil {
    17  		return err
    18  	}
    19  	defer f.Close()
    20  
    21  	err = upload(ctx, f, sendFunc)
    22  	if err != io.EOF {
    23  		return err
    24  	}
    25  	return nil
    26  }
    27  
    28  func upload(ctx context.Context, reader io.Reader, sendFunc func([]byte, error) error) error {
    29  	var err error
    30  	for err == nil {
    31  		buffer := make([]byte, BUFFER_SIZE)
    32  		_, readErr := reader.Read(buffer)
    33  
    34  		err = sendFunc(buffer, readErr)
    35  	}
    36  	return err
    37  }
    38  
    39  func Forward(ctx context.Context, sendFunc func([]byte, error) error, recFunc func() ([]byte, error)) error {
    40  	var err error
    41  	for err == nil {
    42  		var buffer []byte
    43  		buffer, readErr := recFunc()
    44  
    45  		err = sendFunc(buffer, readErr)
    46  	}
    47  	return err
    48  }
    49  
    50  func ReceiveFile(ctx context.Context, filename string, recFunc func() ([]byte, error)) error {
    51  	f, err := os.OpenFile(filename, os.O_CREATE, 0644)
    52  	if err != nil {
    53  		return err
    54  	}
    55  
    56  	err = receive(ctx, f, recFunc)
    57  	if err != io.EOF {
    58  		// if hit an actual error then we want to clean up the file
    59  		f.Close()
    60  		os.Remove(filename)
    61  		return err
    62  	}
    63  	defer f.Close()
    64  
    65  	return nil
    66  }
    67  
    68  func receive(ctx context.Context, writer io.Writer, recFunc func() ([]byte, error)) error {
    69  	var err error
    70  	for err == nil {
    71  		var buffer []byte
    72  		buffer, err = recFunc()
    73  
    74  		_, writeErr := writer.Write(buffer)
    75  		if writeErr != nil {
    76  			return writeErr
    77  		}
    78  	}
    79  	return err
    80  }