github.com/apptainer/singularity@v3.1.1+incompatible/pkg/client/shub/pull.go (about)

     1  // Copyright (c) 2018, Sylabs Inc. All rights reserved.
     2  // This software is licensed under a 3-clause BSD license. Please consult the
     3  // LICENSE.md file distributed with the sources of this project regarding your
     4  // rights to use or distribute this software.
     5  
     6  package client
     7  
     8  import (
     9  	"fmt"
    10  	"io"
    11  	"net/http"
    12  	"os"
    13  	"time"
    14  
    15  	"github.com/sylabs/singularity/internal/pkg/sylog"
    16  	util "github.com/sylabs/singularity/pkg/client/library"
    17  	useragent "github.com/sylabs/singularity/pkg/util/user-agent"
    18  	pb "gopkg.in/cheggaaa/pb.v1"
    19  )
    20  
    21  // Timeout for an image pull in seconds (2 hours)
    22  const pullTimeout = 7200
    23  
    24  // DownloadImage will retrieve an image from the Container Singularityhub,
    25  // saving it into the specified file
    26  func DownloadImage(filePath string, shubRef string, force, noHTTPS bool) (err error) {
    27  	sylog.Debugf("Downloading container from Shub")
    28  
    29  	// use custom parser to make sure we have a valid shub URI
    30  	if ok := isShubPullRef(shubRef); !ok {
    31  		sylog.Fatalf("Invalid shub URI")
    32  	}
    33  
    34  	ShubURI, err := shubParseReference(shubRef)
    35  	if err != nil {
    36  		return fmt.Errorf("Failed to parse shub URI: %v", err)
    37  	}
    38  
    39  	if filePath == "" {
    40  		filePath = fmt.Sprintf("%s_%s.simg", ShubURI.container, ShubURI.tag)
    41  		sylog.Infof("Download filename not provided. Downloading to: %s\n", filePath)
    42  	}
    43  
    44  	if !force {
    45  		if _, err := os.Stat(filePath); err == nil {
    46  			return fmt.Errorf("image file already exists - will not overwrite")
    47  		}
    48  	}
    49  
    50  	// Get the image manifest
    51  	manifest, err := getManifest(ShubURI, noHTTPS)
    52  	if err != nil {
    53  		return fmt.Errorf("Failed to get manifest from Shub: %v", err)
    54  	}
    55  
    56  	// Get the image based on the manifest
    57  	httpc := http.Client{
    58  		Timeout: pullTimeout * time.Second,
    59  	}
    60  
    61  	req, err := http.NewRequest(http.MethodGet, manifest.Image, nil)
    62  	if err != nil {
    63  		return err
    64  	}
    65  	req.Header.Set("User-Agent", useragent.Value())
    66  
    67  	if noHTTPS {
    68  		req.URL.Scheme = "http"
    69  	}
    70  
    71  	// Do the request, if status isn't success, return error
    72  	resp, err := httpc.Do(req)
    73  	if resp == nil {
    74  		return fmt.Errorf("No response received from singularity hub")
    75  	}
    76  	if resp.StatusCode == http.StatusNotFound {
    77  		return fmt.Errorf("The requested image was not found in singularity hub")
    78  	}
    79  	sylog.Debugf("%s response received, beginning image download\n", resp.Status)
    80  
    81  	if resp.StatusCode != http.StatusOK {
    82  		jRes, err := util.ParseErrorBody(resp.Body)
    83  		if err != nil {
    84  			jRes = util.ParseErrorResponse(resp)
    85  		}
    86  		return fmt.Errorf("Download did not succeed: %d %s\n\t%v",
    87  			jRes.Error.Code, jRes.Error.Status, jRes.Error.Message)
    88  	}
    89  
    90  	// Perms are 777 *prior* to umask
    91  	out, err := os.OpenFile(filePath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0777)
    92  	if err != nil {
    93  		return err
    94  	}
    95  	defer out.Close()
    96  
    97  	sylog.Debugf("Created output file: %s\n", filePath)
    98  
    99  	bodySize := resp.ContentLength
   100  	bar := pb.New(int(bodySize)).SetUnits(pb.U_BYTES)
   101  	if sylog.GetLevel() < 0 {
   102  		bar.NotPrint = true
   103  	}
   104  	bar.ShowTimeLeft = true
   105  	bar.ShowSpeed = true
   106  	bar.Start()
   107  
   108  	// create proxy reader
   109  	bodyProgress := bar.NewProxyReader(resp.Body)
   110  
   111  	// Write the body to file
   112  	bytesWritten, err := io.Copy(out, bodyProgress)
   113  	if err != nil {
   114  		return err
   115  	}
   116  	// Simple check to make sure image received is the correct size
   117  	if bytesWritten != resp.ContentLength {
   118  		return fmt.Errorf("Image received is not the right size. Supposed to be: %v  Actually: %v", resp.ContentLength, bytesWritten)
   119  	}
   120  
   121  	bar.Finish()
   122  
   123  	sylog.Debugf("Download complete\n")
   124  
   125  	return err
   126  }