io.go

package simplecsv

import (
	"encoding/csv"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

// SimpleCsv is the type for simple csv
type SimpleCsv [][]string

// CreateEmptyCsv creates an empty CSV with the headers passed as a slice.
// The headers are copied, changes to the original slice don't affect the csv.
// At least one header is required. Header names must be unique, like database
// columns: if columnNames is empty or contains duplicates the returned error
// is non-nil and the returned csv is nil.
func CreateEmptyCsv(columnNames []string) (SimpleCsv, error) {
	if len(columnNames) == 0 {
		return nil, fmt.Errorf("simplecsv: at least one header is required")
	}
	if hasDuplicateHeaders(columnNames) {
		return nil, fmt.Errorf("simplecsv: duplicate header names are not allowed: %v", columnNames)
	}
	a := make([][]string, 1)
	a[0] = copyRow(columnNames)
	return a, nil
}

// CreateEmpyCsv creates an empty CSV with the headers passed as a slice
//
// Deprecated: use CreateEmptyCsv, this function was kept for compatibility
// with the misspelled name
func CreateEmpyCsv(columnNames []string) (SimpleCsv, error) {
	return CreateEmptyCsv(columnNames)
}

// MustCreateEmptyCsv is like CreateEmptyCsv but panics if the headers are
// invalid (empty or duplicated). It is handy in tests and small scripts
// where the header list is part of the program, not a runtime condition:
//
//	people := simplecsv.MustCreateEmptyCsv([]string{"id", "name"})
func MustCreateEmptyCsv(headers []string) SimpleCsv {
	s, err := CreateEmptyCsv(headers)
	if err != nil {
		panic(err)
	}
	return s
}

// ReadCsvFile reads a file and returns a SimpleCsv.
// The second value is false if there's an error reading the file or parsing the CSV.
func ReadCsvFile(filename string) (SimpleCsv, bool) {
	s, err := ReadCsvFileE(filename)
	return s, err == nil
}

// ReadCsvFileE reads a file and returns a SimpleCsv or an error.
func ReadCsvFileE(filename string) (SimpleCsv, error) {
	return ReadCsvFileComma(filename, ',')
}

// ReadCsvFileComma reads a delimited file (for example a semicolon or tab
// separated file) and returns a SimpleCsv or an error.
func ReadCsvFileComma(filename string, comma rune) (SimpleCsv, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, fmt.Errorf("simplecsv: cannot open %q: %w", filename, err)
	}
	defer file.Close()
	return ReadCsv(file, comma)
}

// ReadCsv reads all the records from r using comma as the field separator
// and returns a SimpleCsv or an error.
// The entire input is read into memory; use ReadCsvLimit for bounded reads.
//
// The first record is treated as the header row. Header names must be unique,
// like database columns: a file with duplicate headers in the first row is
// rejected with an error. All rows must have the same number of fields as the
// header row: a ragged file is rejected with an error. This makes uniform row
// width an invariant of a SimpleCsv, alongside the unique-headers invariant.
//
// A UTF-8 byte order mark (BOM, written by Excel and many Windows tools at
// the start of the file) is stripped from the first cell of the first row, so
// the first header is "id" and not "\ufeffid".
func ReadCsv(r io.Reader, comma rune) (SimpleCsv, error) {
	reader := csv.NewReader(r)
	reader.Comma = comma
	// FieldsPerRecord 0 (the encoding/csv default) records the field count of
	// the first record and rejects any subsequent record with a different field
	// count. Set it explicitly: ragged files are rejected, never accepted.
	reader.FieldsPerRecord = 0
	allRecords, err := reader.ReadAll()
	if err != nil {
		return nil, fmt.Errorf("simplecsv: cannot parse csv: %w", err)
	}
	if len(allRecords) == 0 {
		return allRecords, nil
	}
	stripBOM(allRecords)
	if hasDuplicateHeaders(allRecords[0]) {
		return nil, fmt.Errorf("simplecsv: duplicate header names are not allowed: %v", allRecords[0])
	}
	return allRecords, nil
}

// ReadCsvLimit reads records from r using comma as the field separator and
// returns a SimpleCsv or an error. maxRecords is the maximum number of CSV
// records, including the header row, and maxBytes is the maximum number of
// input bytes passed to the CSV parser. A zero limit disables that limit;
// negative limits are rejected. If a limit is exceeded, the returned csv is
// nil.
//
// Unlike ReadCsv, this function does not use csv.Reader.ReadAll. Use positive
// limits when reading input that may be large or attacker-controlled.
//
// Like ReadCsv, a UTF-8 byte order mark (BOM) is stripped from the first
// cell of the first row if present.
func ReadCsvLimit(r io.Reader, comma rune, maxRecords, maxBytes int64) (SimpleCsv, error) {
	if maxRecords < 0 || maxBytes < 0 {
		return nil, fmt.Errorf("simplecsv: read limits must not be negative")
	}

	input := io.Reader(r)
	var byteLimit *io.LimitedReader
	if maxBytes > 0 {
		byteLimit = &io.LimitedReader{R: r, N: maxBytes}
		input = byteLimit
	}

	reader := csv.NewReader(input)
	reader.Comma = comma
	// Keep the same uniform-width invariant as ReadCsv.
	reader.FieldsPerRecord = 0

	records := make(SimpleCsv, 0)
	for {
		if maxRecords > 0 && int64(len(records)) >= maxRecords {
			_, err := reader.Read()
			if err == nil {
				return nil, fmt.Errorf("simplecsv: csv exceeds maximum of %d records", maxRecords)
			}
			if err == io.EOF {
				if exceeded, limitErr := csvByteLimitStatus(byteLimit, r); limitErr != nil {
					return nil, limitErr
				} else if exceeded {
					return nil, fmt.Errorf("simplecsv: csv exceeds maximum of %d bytes", maxBytes)
				}
				break
			}
			if exceeded, limitErr := csvByteLimitStatus(byteLimit, r); limitErr != nil {
				return nil, limitErr
			} else if exceeded {
				return nil, fmt.Errorf("simplecsv: csv exceeds maximum of %d bytes", maxBytes)
			}
			return nil, fmt.Errorf("simplecsv: cannot parse csv: %w", err)
		}

		record, err := reader.Read()
		if err == io.EOF {
			if exceeded, limitErr := csvByteLimitStatus(byteLimit, r); limitErr != nil {
				return nil, limitErr
			} else if exceeded {
				return nil, fmt.Errorf("simplecsv: csv exceeds maximum of %d bytes", maxBytes)
			}
			break
		}
		if err != nil {
			if exceeded, limitErr := csvByteLimitStatus(byteLimit, r); limitErr != nil {
				return nil, limitErr
			} else if exceeded {
				return nil, fmt.Errorf("simplecsv: csv exceeds maximum of %d bytes", maxBytes)
			}
			return nil, fmt.Errorf("simplecsv: cannot parse csv: %w", err)
		}
		records = append(records, record)
	}

	if len(records) == 0 {
		return records, nil
	}
	stripBOM(records)
	if hasDuplicateHeaders(records[0]) {
		return nil, fmt.Errorf("simplecsv: duplicate header names are not allowed: %v", records[0])
	}
	return records, nil
}

// stripBOM removes a UTF-8 byte order mark from the first cell of the first
// row if present. Excel and many Windows tools write one at the start of the
// file; without stripping, the first header becomes "\ufeffid" and every
// name-based lookup silently fails. Only the first cell of the first row is
// touched: a BOM elsewhere in the file is left verbatim.
func stripBOM(records SimpleCsv) {
	if len(records) == 0 || len(records[0]) == 0 {
		return
	}
	records[0][0] = strings.TrimPrefix(records[0][0], "\ufeff")
}

// csvByteLimitStatus checks whether the limited reader stopped at its byte
// limit rather than because the source reached EOF. The one-byte probe is
// necessary because io.LimitedReader cannot distinguish those cases itself.
func csvByteLimitStatus(limit *io.LimitedReader, source io.Reader) (bool, error) {
	if limit == nil || limit.N > 0 {
		return false, nil
	}
	var probe [1]byte
	n, err := source.Read(probe[:])
	if n > 0 {
		return true, nil
	}
	if err == nil || err == io.EOF {
		return false, nil
	}
	return false, fmt.Errorf("simplecsv: cannot read csv: %w", err)
}

// WriteCsvFile writes the SimpleCsv to a file.
// Returns false if there's an error creating or writing the file.
func (s SimpleCsv) WriteCsvFile(filename string) bool {
	return s.WriteCsvFileE(filename) == nil
}

// WriteCsvFileE writes the SimpleCsv to a file and returns an error on failure.
func (s SimpleCsv) WriteCsvFileE(filename string) error {
	return s.WriteCsvFileComma(filename, ',')
}

// WriteCsvFileComma writes the SimpleCsv to a file using comma as the field
// separator and returns an error on failure.
//
// The write is atomic: the csv is first written to a temporary file in the
// same directory as filename and then renamed over it, so a failed or
// interrupted write cannot truncate or corrupt the destination. The temporary
// file (and therefore the destination after the rename) is created with mode
// 0600 (owner-readable only), so written files are not world-readable by
// default; callers that need different permissions can os.Chmod the result.
// Because the destination is replaced with os.Rename rather than opened, a
// symlink planted at filename is replaced instead of followed: the write does
// not follow a symlink at the destination.
func (s SimpleCsv) WriteCsvFileComma(filename string, comma rune) error {
	dir := filepath.Dir(filename)
	tmp, err := os.CreateTemp(dir, ".simplecsv-*")
	if err != nil {
		return fmt.Errorf("simplecsv: cannot create temp file for %q: %w", filename, err)
	}
	tmpName := tmp.Name()
	// os.CreateTemp opens with 0600 on Unix; chmod explicitly so the final
	// file is owner-readable only by default regardless of CreateTemp's mode.
	defer os.Remove(tmpName) // no-op once the temp has been renamed
	if cerr := os.Chmod(tmpName, 0600); cerr != nil {
		tmp.Close()
		return fmt.Errorf("simplecsv: cannot set permissions on temp file for %q: %w", filename, cerr)
	}
	werr := s.WriteTo(tmp, comma)
	cerr := tmp.Close()
	if werr != nil {
		return fmt.Errorf("simplecsv: cannot write %q: %w", filename, werr)
	}
	if cerr != nil {
		return fmt.Errorf("simplecsv: cannot close temp file for %q: %w", filename, cerr)
	}
	if rerr := os.Rename(tmpName, filename); rerr != nil {
		return fmt.Errorf("simplecsv: cannot rename temp file to %q: %w", filename, rerr)
	}
	return nil
}

// WriteTo writes the SimpleCsv to w using comma as the field separator
// and returns an error on failure.
func (s SimpleCsv) WriteTo(w io.Writer, comma rune) error {
	cw := csv.NewWriter(w)
	cw.Comma = comma
	if err := cw.WriteAll(s); err != nil {
		return fmt.Errorf("simplecsv: cannot write csv: %w", err)
	}
	return nil
}