github.com/blend/go-sdk@v1.20220411.3/sh/limit_bytes.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package sh
     9  
    10  import (
    11  	"io"
    12  	"sync"
    13  
    14  	"github.com/blend/go-sdk/ex"
    15  )
    16  
    17  // Errors
    18  const (
    19  	ErrMaxBytesWriterCapacityLimit ex.Class = "write failed; maximum capacity reached or would be exceede"
    20  )
    21  
    22  // LimitBytes returns a new max bytes writer.
    23  func LimitBytes(maxBytes int, inner io.Writer) *MaxBytesWriter {
    24  	return &MaxBytesWriter{
    25  		max:   maxBytes,
    26  		inner: inner,
    27  	}
    28  }
    29  
    30  // MaxBytesWriter returns a maximum bytes writer.
    31  type MaxBytesWriter struct {
    32  	sync.Mutex
    33  
    34  	max   int
    35  	count int
    36  	inner io.Writer
    37  }
    38  
    39  // Max returns the maximum bytes allowed.
    40  func (mbw *MaxBytesWriter) Max() int {
    41  	return int(mbw.max)
    42  }
    43  
    44  // Count returns the current written bytes..
    45  func (mbw *MaxBytesWriter) Count() int {
    46  	return int(mbw.count)
    47  }
    48  
    49  func (mbw *MaxBytesWriter) Write(contents []byte) (int, error) {
    50  	mbw.Lock()
    51  	defer mbw.Unlock()
    52  
    53  	if mbw.count >= mbw.max {
    54  		return 0, ex.New(ErrMaxBytesWriterCapacityLimit)
    55  	}
    56  	if len(contents)+mbw.count >= mbw.max {
    57  		return 0, ex.New(ErrMaxBytesWriterCapacityLimit)
    58  	}
    59  
    60  	written, err := mbw.inner.Write(contents)
    61  	mbw.count += written
    62  	if err != nil {
    63  		return written, ex.New(err)
    64  	}
    65  	return written, nil
    66  }