github.com/opcr-io/oras-go/v2@v2.0.0-20231122155130-eb4260d8a0ae/internal/registryutil/proxy.go (about)

     1  /*
     2  Copyright The ORAS Authors.
     3  Licensed under the Apache License, Version 2.0 (the "License");
     4  you may not use this file except in compliance with the License.
     5  You may obtain a copy of the License at
     6  
     7  http://www.apache.org/licenses/LICENSE-2.0
     8  
     9  Unless required by applicable law or agreed to in writing, software
    10  distributed under the License is distributed on an "AS IS" BASIS,
    11  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  See the License for the specific language governing permissions and
    13  limitations under the License.
    14  */
    15  
    16  package registryutil
    17  
    18  import (
    19  	"context"
    20  	"io"
    21  	"sync"
    22  
    23  	"github.com/opcr-io/oras-go/v2/content"
    24  	"github.com/opcr-io/oras-go/v2/internal/cas"
    25  	"github.com/opcr-io/oras-go/v2/internal/ioutil"
    26  	"github.com/opcr-io/oras-go/v2/registry"
    27  	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
    28  )
    29  
    30  // ReferenceStorage represents a CAS that supports registry.ReferenceFetcher.
    31  type ReferenceStorage interface {
    32  	content.ReadOnlyStorage
    33  	registry.ReferenceFetcher
    34  }
    35  
    36  // Proxy is a caching proxy dedicated for registry.ReferenceFetcher.
    37  // The first fetch call of a described content will read from the remote and
    38  // cache the fetched content.
    39  // The subsequent fetch call will read from the local cache.
    40  type Proxy struct {
    41  	registry.ReferenceFetcher
    42  	*cas.Proxy
    43  }
    44  
    45  // NewProxy creates a proxy for the `base` ReferenceStorage, using the `cache`
    46  // storage as the cache.
    47  func NewProxy(base ReferenceStorage, cache content.Storage) *Proxy {
    48  	return &Proxy{
    49  		ReferenceFetcher: base,
    50  		Proxy:            cas.NewProxy(base, cache),
    51  	}
    52  }
    53  
    54  // FetchReference fetches the content identified by the reference from the
    55  // remote and cache the fetched content.
    56  func (p *Proxy) FetchReference(ctx context.Context, reference string) (ocispec.Descriptor, io.ReadCloser, error) {
    57  	target, rc, err := p.ReferenceFetcher.FetchReference(ctx, reference)
    58  	if err != nil {
    59  		return ocispec.Descriptor{}, nil, err
    60  	}
    61  
    62  	// skip caching if the content already exists in cache
    63  	exists, err := p.Cache.Exists(ctx, target)
    64  	if err != nil {
    65  		return ocispec.Descriptor{}, nil, err
    66  	}
    67  	if exists {
    68  		return target, rc, nil
    69  	}
    70  
    71  	// cache content while reading
    72  	pr, pw := io.Pipe()
    73  	var wg sync.WaitGroup
    74  	wg.Add(1)
    75  	var pushErr error
    76  	go func() {
    77  		defer wg.Done()
    78  		pushErr = p.Cache.Push(ctx, target, pr)
    79  		if pushErr != nil {
    80  			pr.CloseWithError(pushErr)
    81  		}
    82  	}()
    83  	closer := ioutil.CloserFunc(func() error {
    84  		rcErr := rc.Close()
    85  		if err := pw.Close(); err != nil {
    86  			return err
    87  		}
    88  		wg.Wait()
    89  		if pushErr != nil {
    90  			return pushErr
    91  		}
    92  		return rcErr
    93  	})
    94  
    95  	return target, struct {
    96  		io.Reader
    97  		io.Closer
    98  	}{
    99  		Reader: io.TeeReader(rc, pw),
   100  		Closer: closer,
   101  	}, nil
   102  }