github.com/castai/kvisor@v1.7.1-0.20240516114728-b3572a2607b5/pkg/blobscache/client.go (about)

     1  package blobscache
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"fmt"
     7  	"io"
     8  	"net/http"
     9  	"time"
    10  
    11  	ia "github.com/castai/image-analyzer"
    12  	json "github.com/json-iterator/go"
    13  )
    14  
    15  func NewRemoteBlobsCacheClient(serverURL string) ia.CacheClient {
    16  	return &remoteBlobsCache{
    17  		url: serverURL,
    18  		client: &http.Client{
    19  			Timeout: 10 * time.Second,
    20  		},
    21  	}
    22  }
    23  
    24  type remoteBlobsCache struct {
    25  	url    string
    26  	client *http.Client
    27  }
    28  
    29  func (c *remoteBlobsCache) PutBlob(ctx context.Context, key string, blob []byte) error {
    30  	var buf bytes.Buffer
    31  	if err := json.NewEncoder(&buf).Encode(&PubBlobRequest{
    32  		Key:  key,
    33  		Blob: blob,
    34  	}); err != nil {
    35  		return err
    36  	}
    37  	req, err := http.NewRequestWithContext(ctx, "POST", fmt.Sprintf("%s/blobscache/PutBlob", c.url), &buf)
    38  	if err != nil {
    39  		return err
    40  	}
    41  	req.Header.Add("Content-Type", "application/json")
    42  	resp, err := c.client.Do(req)
    43  	if err != nil {
    44  		return err
    45  	}
    46  	defer resp.Body.Close()
    47  	if st := resp.StatusCode; st != http.StatusOK {
    48  		errMsg, _ := io.ReadAll(resp.Body)
    49  		return fmt.Errorf("put blob failed, response status=%d: %v", st, string(errMsg))
    50  	}
    51  	return nil
    52  }
    53  
    54  func (c *remoteBlobsCache) GetBlob(ctx context.Context, key string) ([]byte, error) {
    55  	var buf bytes.Buffer
    56  	if err := json.NewEncoder(&buf).Encode(&GetBlobRequest{
    57  		Key: key,
    58  	}); err != nil {
    59  		return nil, err
    60  	}
    61  	req, err := http.NewRequestWithContext(ctx, "POST", fmt.Sprintf("%s/blobscache/GetBlob", c.url), &buf)
    62  	if err != nil {
    63  		return nil, err
    64  	}
    65  	req.Header.Add("Content-Type", "application/json")
    66  	resp, err := c.client.Do(req)
    67  	if err != nil {
    68  		return nil, err
    69  	}
    70  	defer resp.Body.Close()
    71  	if st := resp.StatusCode; st != http.StatusOK {
    72  		if st == http.StatusNotFound {
    73  			return nil, ia.ErrCacheNotFound
    74  		}
    75  		errMsg, _ := io.ReadAll(resp.Body)
    76  		return nil, fmt.Errorf("put blob failed, response status=%d: %v", st, string(errMsg))
    77  	}
    78  
    79  	var res GetBlobResponse
    80  	if err := json.NewDecoder(resp.Body).Decode(&res); err != nil {
    81  		return nil, err
    82  	}
    83  	return res.Blob, nil
    84  }
    85  
    86  func NewMockBlobsCacheClient() ia.CacheClient {
    87  	return &mockBlobsCacheClient{}
    88  }
    89  
    90  type mockBlobsCacheClient struct {
    91  }
    92  
    93  func (m mockBlobsCacheClient) PutBlob(ctx context.Context, key string, blob []byte) error {
    94  	return nil
    95  }
    96  
    97  func (m mockBlobsCacheClient) GetBlob(ctx context.Context, key string) ([]byte, error) {
    98  	return nil, ia.ErrCacheNotFound
    99  }