github.com/web-platform-tests/wpt.fyi@v0.0.0-20240530210107-70cf978996f1/api/screenshot/cache.go (about)

     1  // Copyright 2019 The WPT Dashboard Project. All rights reserved.
     2  // Use of this source code is governed by a BSD-style license that can be
     3  // found in the LICENSE file.
     4  
     5  package screenshot
     6  
     7  import (
     8  	"context"
     9  	"encoding/json"
    10  	"errors"
    11  	"io"
    12  	"net/http"
    13  	"sync"
    14  
    15  	"cloud.google.com/go/storage"
    16  
    17  	"github.com/web-platform-tests/wpt.fyi/api/receiver"
    18  	"github.com/web-platform-tests/wpt.fyi/shared"
    19  )
    20  
    21  func parseParams(r *http.Request) (browser, browserVersion, os, osVersion string) {
    22  	browser = r.FormValue("browser")
    23  	browserVersion = r.FormValue("browser_version")
    24  	os = r.FormValue("os")
    25  	osVersion = r.FormValue("os_version")
    26  
    27  	return browser, browserVersion, os, osVersion
    28  }
    29  
    30  func getHashesHandler(w http.ResponseWriter, r *http.Request) {
    31  	ctx := r.Context()
    32  	ds := shared.NewAppEngineDatastore(ctx, false)
    33  
    34  	browser, browserVersion, os, osVersion := parseParams(r)
    35  	hashes, err := RecentScreenshotHashes(ds, browser, browserVersion, os, osVersion, nil)
    36  	if err != nil {
    37  		http.Error(w, err.Error(), http.StatusInternalServerError)
    38  
    39  		return
    40  	}
    41  	response, err := json.Marshal(hashes)
    42  	if err != nil {
    43  		http.Error(w, err.Error(), http.StatusInternalServerError)
    44  
    45  		return
    46  	}
    47  	_, err = w.Write(response)
    48  	if err != nil {
    49  		logger := shared.GetLogger(ctx)
    50  		logger.Warningf("Failed to write data in api/screenshots/hashes handler: %s", err.Error())
    51  	}
    52  }
    53  
    54  func uploadScreenshotHandler(w http.ResponseWriter, r *http.Request) {
    55  	if r.Method != http.MethodPost {
    56  		http.Error(w, "Only POST is supported", http.StatusMethodNotAllowed)
    57  
    58  		return
    59  	}
    60  
    61  	ctx := r.Context()
    62  	aeAPI := shared.NewAppEngineAPI(ctx)
    63  	if receiver.AuthenticateUploader(aeAPI, r) != receiver.InternalUsername {
    64  		http.Error(w, "This is a private API.", http.StatusUnauthorized)
    65  
    66  		return
    67  	}
    68  
    69  	// nolint:godox // TODO(Hexcles): Abstract and mock the GCS utilities in shared.
    70  	gcs, err := storage.NewClient(ctx)
    71  	if err != nil {
    72  		http.Error(w, err.Error(), http.StatusInternalServerError)
    73  
    74  		return
    75  	}
    76  	/* #nosec G101 */
    77  	bucketName := "wptd-screenshots-staging"
    78  	if aeAPI.GetHostname() == "wpt.fyi" {
    79  		/* #nosec G101 */
    80  		bucketName = "wptd-screenshots"
    81  	}
    82  	bucket := gcs.Bucket(bucketName)
    83  
    84  	browser, browserVersion, os, osVersion := parseParams(r)
    85  	hashMethod := r.FormValue("hash_method")
    86  	if r.MultipartForm == nil || r.MultipartForm.File == nil || len(r.MultipartForm.File["screenshot"]) == 0 {
    87  		http.Error(w, "no screenshot file found", http.StatusBadRequest)
    88  
    89  		return
    90  	}
    91  
    92  	fhs := r.MultipartForm.File["screenshot"]
    93  	errors := make(chan error, len(fhs))
    94  	var wg sync.WaitGroup
    95  	wg.Add(len(fhs))
    96  	for i := range fhs {
    97  		go func(i int) {
    98  			defer wg.Done()
    99  			f, err := fhs[i].Open()
   100  			if err != nil {
   101  				errors <- err
   102  
   103  				return
   104  			}
   105  			defer f.Close()
   106  			if err := storeScreenshot(ctx, bucket, hashMethod, browser, browserVersion, os, osVersion, f); err != nil {
   107  				errors <- err
   108  			}
   109  		}(i)
   110  	}
   111  	wg.Wait()
   112  	close(errors)
   113  
   114  	me := shared.NewMultiErrorFromChan(errors, "storing screenshots to GCS")
   115  	if me != nil {
   116  		http.Error(w, me.Error(), http.StatusInternalServerError)
   117  
   118  		return
   119  	}
   120  	w.WriteHeader(http.StatusCreated)
   121  }
   122  
   123  func storeScreenshot(
   124  	ctx context.Context,
   125  	bucket *storage.BucketHandle,
   126  	hashMethod,
   127  	browser,
   128  	browserVersion,
   129  	os,
   130  	osVersion string,
   131  	f io.ReadSeeker,
   132  ) error {
   133  	if hashMethod == "" {
   134  		hashMethod = "sha1"
   135  	}
   136  	s := NewScreenshot(browser, browserVersion, os, osVersion)
   137  	if err := s.SetHashFromFile(f, hashMethod); err != nil {
   138  		return err
   139  	}
   140  	// Need to reset the file after hashing it.
   141  	if _, err := f.Seek(0, io.SeekStart); err != nil {
   142  		logger := shared.GetLogger(ctx)
   143  		logger.Warningf("Failed to reset file: %s", err.Error())
   144  	}
   145  
   146  	ds := shared.NewAppEngineDatastore(ctx, false)
   147  
   148  	o := bucket.Object(s.Hash() + ".png")
   149  	if _, err := o.Attrs(ctx); errors.Is(err, storage.ErrObjectNotExist) {
   150  		w := o.NewWriter(ctx)
   151  		// Screenshots are small; disable chunking for better performance.
   152  		w.ChunkSize = 0
   153  		if _, err := io.Copy(w, f); err != nil {
   154  			return err
   155  		}
   156  		if err := w.Close(); err != nil {
   157  			return err
   158  		}
   159  	}
   160  
   161  	// Write to Datastore last.
   162  	return s.Store(ds)
   163  }