MyGoNavi/internal/app/sql_import_source.go

package app

import (
	"bufio"
	"bytes"
	"compress/gzip"
	"fmt"
	"io"
	"os"
	"strings"

	"golang.org/x/text/encoding/unicode"
	"golang.org/x/text/transform"
)

// SQLImportSourceOptions defines safety limits applied while decoding an SQL
// import source. Zero values select production-safe defaults.
type SQLImportSourceOptions struct {
	MaxDecodedBytes            int64
	MaxCompressionRatio        float64
	MinCompressedBytesForRatio int64
	// RawObserver receives the original on-disk bytes as they are read. It is
	// suitable for streaming hashes and never receives decoded SQL content.
	RawObserver io.Writer
}

type SQLImportSourceLimitKind string

const (
	SQLImportSourceDecodedByteLimit            SQLImportSourceLimitKind = "decoded_byte_limit"
	SQLImportSourceCompressionRatio            SQLImportSourceLimitKind = "compression_ratio"
	defaultSQLImportMaxDecodedBytes                                     = int64(16 << 30)
	defaultSQLImportMaxCompressionRatio                                 = 1000.0
	defaultSQLImportMinCompressedBytesForRatio                          = int64(1)
)

// SQLImportSourceLimitError reports a streaming safety limit without exposing
// source content.
type SQLImportSourceLimitError struct {
	Kind                SQLImportSourceLimitKind
	Limit               int64
	DecodedBytes        int64
	CompressedBytes     int64
	Ratio               float64
	MaxCompressionRatio float64
}

func (err *SQLImportSourceLimitError) Error() string {
	return fmt.Sprintf("SQL import source exceeded %s safety limit", err.Kind)
}

type sqlImportDecodedLimitReader struct {
	reader                  io.Reader
	maxBytes                int64
	decodedBytes            int64
	compressed              *sqlImportCountingReader
	maxCompressionRatio     float64
	minCompressedRatioBytes int64
}

func (reader *sqlImportDecodedLimitReader) Read(buffer []byte) (int, error) {
	remaining := reader.maxBytes - reader.decodedBytes
	if remaining < 0 {
		remaining = 0
	}
	probeSize := int64(len(buffer))
	if probeSize > remaining+1 {
		probeSize = remaining + 1
	}
	n, readErr := reader.reader.Read(buffer[:int(probeSize)])
	projectedDecodedBytes := reader.decodedBytes + int64(n)
	if int64(n) > remaining {
		allowed := int(remaining)
		reader.decodedBytes += int64(allowed)
		return allowed, &SQLImportSourceLimitError{
			Kind:         SQLImportSourceDecodedByteLimit,
			Limit:        reader.maxBytes,
			DecodedBytes: reader.decodedBytes,
		}
	}
	if reader.compressed != nil && reader.compressed.bytes >= reader.minCompressedRatioBytes {
		ratio := float64(projectedDecodedBytes) / float64(reader.compressed.bytes)
		if ratio > reader.maxCompressionRatio {
			reader.decodedBytes = projectedDecodedBytes
			return 0, &SQLImportSourceLimitError{
				Kind:                SQLImportSourceCompressionRatio,
				DecodedBytes:        projectedDecodedBytes,
				CompressedBytes:     reader.compressed.bytes,
				Ratio:               ratio,
				MaxCompressionRatio: reader.maxCompressionRatio,
			}
		}
	}
	reader.decodedBytes = projectedDecodedBytes
	return n, readErr
}

type sqlImportCountingReader struct {
	reader io.Reader
	bytes  int64
}

func (reader *sqlImportCountingReader) Read(buffer []byte) (int, error) {
	n, err := reader.reader.Read(buffer)
	reader.bytes += int64(n)
	return n, err
}

// SQLImportSource is a streaming, UTF-8 SQL source. Callers must close it.
type SQLImportSource struct {
	io.Reader
	Encoding   string
	Compressed bool
	rawCounter *sqlImportCountingReader
	close      func() error
}

// RawBytesRead returns bytes consumed from the original file. For .sql.gz it
// therefore reports compressed bytes rather than decoded SQL bytes.
func (source *SQLImportSource) RawBytesRead() int64 {
	if source == nil || source.rawCounter == nil {
		return 0
	}
	return source.rawCounter.bytes
}

func (source *SQLImportSource) Close() error {
	if source == nil || source.close == nil {
		return nil
	}
	return source.close()
}

// OpenSQLImportSource opens path without loading it into memory.
func OpenSQLImportSource(path string, options SQLImportSourceOptions) (*SQLImportSource, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	rawReader := io.Reader(file)
	if options.RawObserver != nil {
		rawReader = io.TeeReader(rawReader, options.RawObserver)
	}
	rawCounter := &sqlImportCountingReader{reader: rawReader}
	payload := io.Reader(rawCounter)
	compressed := strings.HasSuffix(strings.ToLower(path), ".sql.gz")
	var compressionCounter *sqlImportCountingReader
	closeSource := file.Close
	if compressed {
		compressionCounter = rawCounter
		gzipReader, gzipErr := gzip.NewReader(rawCounter)
		if gzipErr != nil {
			_ = file.Close()
			return nil, gzipErr
		}
		payload = gzipReader
		closeSource = func() error {
			gzipCloseErr := gzipReader.Close()
			fileCloseErr := file.Close()
			if gzipCloseErr != nil {
				return gzipCloseErr
			}
			return fileCloseErr
		}
	}
	buffered := bufio.NewReader(payload)
	reader := io.Reader(buffered)
	encodingName := "utf-8"
	if prefix, _ := buffered.Peek(3); bytes.Equal(prefix, []byte{0xef, 0xbb, 0xbf}) {
		_, _ = buffered.Discard(3)
	} else if prefix, _ := buffered.Peek(2); bytes.Equal(prefix, []byte{0xff, 0xfe}) {
		_, _ = buffered.Discard(2)
		reader = transform.NewReader(buffered, unicode.UTF16(unicode.LittleEndian, unicode.IgnoreBOM).NewDecoder())
		encodingName = "utf-16le"
	} else if prefix, _ := buffered.Peek(2); bytes.Equal(prefix, []byte{0xfe, 0xff}) {
		_, _ = buffered.Discard(2)
		reader = transform.NewReader(buffered, unicode.UTF16(unicode.BigEndian, unicode.IgnoreBOM).NewDecoder())
		encodingName = "utf-16be"
	}
	maxDecodedBytes := options.MaxDecodedBytes
	if maxDecodedBytes <= 0 {
		maxDecodedBytes = defaultSQLImportMaxDecodedBytes
	}
	maxCompressionRatio := options.MaxCompressionRatio
	if maxCompressionRatio <= 0 {
		maxCompressionRatio = defaultSQLImportMaxCompressionRatio
	}
	minCompressedRatioBytes := options.MinCompressedBytesForRatio
	if minCompressedRatioBytes <= 0 {
		minCompressedRatioBytes = defaultSQLImportMinCompressedBytesForRatio
	}
	reader = &sqlImportDecodedLimitReader{
		reader:                  reader,
		maxBytes:                maxDecodedBytes,
		compressed:              compressionCounter,
		maxCompressionRatio:     maxCompressionRatio,
		minCompressedRatioBytes: minCompressedRatioBytes,
	}
	return &SQLImportSource{
		Reader:     reader,
		Encoding:   encodingName,
		Compressed: compressed,
		rawCounter: rawCounter,
		close:      closeSource,
	}, nil
}