file_size.go 625 B

1234567891011121314151617181920212223242526272829303132333435363738
  1. package security
  2. import (
  3. "io"
  4. )
  5. type hardLimitReader struct {
  6. r io.Reader
  7. left int
  8. }
  9. func (lr *hardLimitReader) Read(p []byte) (n int, err error) {
  10. if lr.left <= 0 {
  11. return 0, newFileSizeError()
  12. }
  13. if len(p) > lr.left {
  14. p = p[0:lr.left]
  15. }
  16. n, err = lr.r.Read(p)
  17. lr.left -= n
  18. return
  19. }
  20. func CheckFileSize(size int, opts Options) error {
  21. if opts.MaxSrcFileSize > 0 && size > opts.MaxSrcFileSize {
  22. return newFileSizeError()
  23. }
  24. return nil
  25. }
  26. func LimitFileSize(r io.Reader, opts Options) io.Reader {
  27. if opts.MaxSrcFileSize > 0 {
  28. return &hardLimitReader{r: r, left: opts.MaxSrcFileSize}
  29. }
  30. return r
  31. }