github.com/nilium/gitlab-runner@v12.5.0+incompatible/commands/helpers/artifacts_test.go (about)

     1  package helpers
     2  
     3  import (
     4  	"archive/zip"
     5  	"bytes"
     6  	"compress/gzip"
     7  	"io"
     8  	"io/ioutil"
     9  	"os"
    10  	"testing"
    11  
    12  	"github.com/sirupsen/logrus"
    13  	"github.com/stretchr/testify/require"
    14  
    15  	"gitlab.com/gitlab-org/gitlab-runner/common"
    16  )
    17  
    18  const (
    19  	artifactsTestArchivedFile  = "archive_file"
    20  	artifactsTestArchivedFile2 = "archive_file2"
    21  )
    22  
    23  type testNetwork struct {
    24  	common.MockNetwork
    25  	downloadState  common.DownloadState
    26  	downloadCalled int
    27  	uploadState    common.UploadState
    28  	uploadCalled   int
    29  	uploadFormat   common.ArtifactFormat
    30  	uploadName     string
    31  	uploadType     string
    32  	uploadedFiles  []string
    33  }
    34  
    35  func (m *testNetwork) DownloadArtifacts(config common.JobCredentials, artifactsFile string) common.DownloadState {
    36  	m.downloadCalled++
    37  
    38  	if m.downloadState == common.DownloadSucceeded {
    39  		file, err := os.Create(artifactsFile)
    40  		if err != nil {
    41  			logrus.Warningln(err)
    42  			return common.DownloadFailed
    43  		}
    44  		defer file.Close()
    45  
    46  		archive := zip.NewWriter(file)
    47  		archive.Create(artifactsTestArchivedFile)
    48  		archive.Close()
    49  	}
    50  	return m.downloadState
    51  }
    52  
    53  func (m *testNetwork) consumeZipUpload(config common.JobCredentials, reader io.Reader, options common.ArtifactsOptions) common.UploadState {
    54  	var buffer bytes.Buffer
    55  	io.Copy(&buffer, reader)
    56  	archive, err := zip.NewReader(bytes.NewReader(buffer.Bytes()), int64(buffer.Len()))
    57  	if err != nil {
    58  		logrus.Warningln(err)
    59  		return common.UploadForbidden
    60  	}
    61  
    62  	for _, file := range archive.File {
    63  		m.uploadedFiles = append(m.uploadedFiles, file.Name)
    64  	}
    65  
    66  	m.uploadFormat = common.ArtifactFormatZip
    67  
    68  	return m.uploadState
    69  }
    70  
    71  func (m *testNetwork) consumeGzipUpload(config common.JobCredentials, reader io.Reader, options common.ArtifactsOptions) common.UploadState {
    72  	var buffer bytes.Buffer
    73  	io.Copy(&buffer, reader)
    74  
    75  	stream := bytes.NewReader(buffer.Bytes())
    76  
    77  	gz, err := gzip.NewReader(stream)
    78  	gz.Multistream(false)
    79  	if err != nil {
    80  		logrus.Warningln("Invalid gzip stream")
    81  		return common.UploadForbidden
    82  	}
    83  
    84  	// Read multiple streams
    85  	for {
    86  		_, err = io.Copy(ioutil.Discard, gz)
    87  		if err != nil {
    88  			logrus.Warningln("Invalid gzip stream")
    89  			return common.UploadForbidden
    90  		}
    91  
    92  		m.uploadedFiles = append(m.uploadedFiles, gz.Name)
    93  
    94  		if gz.Reset(stream) == io.EOF {
    95  			break
    96  		}
    97  		gz.Multistream(false)
    98  	}
    99  
   100  	m.uploadFormat = common.ArtifactFormatGzip
   101  
   102  	return m.uploadState
   103  }
   104  
   105  func (m *testNetwork) consumeRawUpload(config common.JobCredentials, reader io.Reader, options common.ArtifactsOptions) common.UploadState {
   106  	io.Copy(ioutil.Discard, reader)
   107  
   108  	m.uploadedFiles = append(m.uploadedFiles, "raw")
   109  	m.uploadFormat = common.ArtifactFormatRaw
   110  	return m.uploadState
   111  }
   112  
   113  func (m *testNetwork) UploadRawArtifacts(config common.JobCredentials, reader io.Reader, options common.ArtifactsOptions) common.UploadState {
   114  	m.uploadCalled++
   115  
   116  	if m.uploadState == common.UploadSucceeded {
   117  		m.uploadType = options.Type
   118  		m.uploadName = options.BaseName
   119  
   120  		switch options.Format {
   121  		case common.ArtifactFormatZip, common.ArtifactFormatDefault:
   122  			return m.consumeZipUpload(config, reader, options)
   123  
   124  		case common.ArtifactFormatGzip:
   125  			return m.consumeGzipUpload(config, reader, options)
   126  
   127  		case common.ArtifactFormatRaw:
   128  			return m.consumeRawUpload(config, reader, options)
   129  
   130  		default:
   131  			return common.UploadForbidden
   132  		}
   133  	}
   134  
   135  	return m.uploadState
   136  }
   137  
   138  func writeTestFile(t *testing.T, fileName string) {
   139  	err := ioutil.WriteFile(fileName, nil, 0600)
   140  	require.NoError(t, err, "Writing file:", fileName)
   141  }