github.com/google/syzkaller@v0.0.0-20251211124644-a066d2bc4b02/syz-cluster/pkg/blob/storage.go (about)

     1  // Copyright 2024 syzkaller project authors. All rights reserved.
     2  // Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
     3  
     4  package blob
     5  
     6  import (
     7  	"encoding/base64"
     8  	"fmt"
     9  	"io"
    10  	"os"
    11  	"path/filepath"
    12  	"strings"
    13  )
    14  
    15  // Storage is not assumed to be used for partciularly large objects (e.g. GB of size),
    16  // but rather for blobs that risk overwhelming Spanner column size limits.
    17  type Storage interface {
    18  	// Write stores the object uniquely identified by a set of IDs (parts).
    19  	// If it already exists, it will be overwritten.
    20  	// The first argument is the URI which can be used to later retrieve it with Read.
    21  	Write(source io.Reader, parts ...string) (string, error)
    22  	Read(uri string) (io.ReadCloser, error)
    23  }
    24  
    25  var _ Storage = (*LocalStorage)(nil)
    26  
    27  // LocalStorage keeps objets in the specified local directory.
    28  // It's intended to be used only for unit tests.
    29  type LocalStorage struct {
    30  	baseFolder string
    31  }
    32  
    33  func NewLocalStorage(baseFolder string) *LocalStorage {
    34  	return &LocalStorage{baseFolder: baseFolder}
    35  }
    36  
    37  const localStoragePrefix = "local://"
    38  
    39  func (ls *LocalStorage) Write(source io.Reader, parts ...string) (string, error) {
    40  	// A whatever approach that can handle arbitrary inputs.
    41  	name := base64.StdEncoding.EncodeToString([]byte(filepath.Join(parts...)))
    42  	file, err := os.Create(filepath.Join(ls.baseFolder, name))
    43  	if err != nil {
    44  		return "", err
    45  	}
    46  	defer file.Close()
    47  	_, err = io.Copy(file, source)
    48  	if err != nil {
    49  		return "", fmt.Errorf("failed to save data: %w", err)
    50  	}
    51  	return localStoragePrefix + name, nil
    52  }
    53  
    54  func (ls *LocalStorage) Read(uri string) (io.ReadCloser, error) {
    55  	if !strings.HasPrefix(uri, localStoragePrefix) {
    56  		return nil, fmt.Errorf("unsupported URI type")
    57  	}
    58  	// TODO: add some other URI validation checks?
    59  	path := filepath.Join(ls.baseFolder, strings.TrimPrefix(uri, localStoragePrefix))
    60  	return os.Open(path)
    61  }
    62  
    63  func ReadAllBytes(storage Storage, uri string) ([]byte, error) {
    64  	if uri == "" {
    65  		return nil, nil
    66  	}
    67  	reader, err := storage.Read(uri)
    68  	if err != nil {
    69  		return nil, err
    70  	}
    71  	defer reader.Close()
    72  	return io.ReadAll(reader)
    73  }