sponge/main.go

159 lines
3.0 KiB
Go

// sponge - Soaks up all input from stdin and writes it to a file or stdout
// https://gitlab.com/tslocum/sponge
//
// Originally written in C by Tollef Fog Heen
// https://joeyh.name/code/moreutils
//
// Ported by Trevor Slocum
// https://rocketnine.space
package main
import (
"flag"
"fmt"
"io"
"io/ioutil"
"log"
"os"
"os/exec"
)
// BufferSize determines the size of the input/output buffer. The buffer is
// used directly when soaking initially. If the input is larger than the buffer
// it is written to a temporary file instead. The buffer is reused when
// writing to and reading from the temporary file.
const BufferSize = 64 * 1024 // 64 KiB
var (
buf = make([]byte, BufferSize)
bufInitialRead int
appendFile bool
tmpFile *os.File
)
func init() {
log.SetFlags(0)
log.SetPrefix("sponge: ")
}
func main() {
flag.BoolVar(&appendFile, "a", false, "append")
flag.Usage = func() {
fmt.Fprintf(os.Stderr, "%s [-a] [file]: soak up all input from stdin and write it to [file] or stdout\n", os.Args[0])
}
flag.Parse()
err := soak()
if err != nil {
log.Fatal(err)
}
defer cleanUp()
err = squeeze(flag.Arg(0))
if err != nil {
log.Fatal(err)
}
}
func soak() error {
var err error
bufInitialRead, err = io.ReadFull(os.Stdin, buf)
if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
return err
} else if bufInitialRead == BufferSize && err != io.ErrUnexpectedEOF {
tmpFile, err = ioutil.TempFile("", "sponge-*.tmp")
if err != nil {
return err
}
_, err = tmpFile.Write(buf[:bufInitialRead])
if err != nil {
tmpFile.Close()
return err
}
_, err = io.CopyBuffer(tmpFile, os.Stdin, buf)
if err != nil {
tmpFile.Close()
return err
}
_, err = tmpFile.Seek(io.SeekStart, 0)
if err != nil {
tmpFile.Close()
return err
}
}
return nil
}
func squeeze(filename string) error {
var out io.Writer
if filename != "" {
regularFile := true
fileInfo, err := os.Stat(filename)
if err != nil {
if !os.IsNotExist(err) {
return err
}
} else {
regularFile = fileInfo.Mode().IsRegular()
}
if tmpFile != nil && regularFile && !appendFile {
tmpFile.Close()
// Rename temporary file
err = os.Rename(tmpFile.Name(), filename)
if err != nil {
// Fall back to mv
cmd := exec.Command("mv", tmpFile.Name(), filename)
err = cmd.Run()
}
tmpFile = nil
return err
}
openFlags := os.O_WRONLY
if regularFile {
openFlags |= os.O_CREATE
}
if appendFile {
openFlags |= os.O_APPEND
} else if regularFile {
openFlags |= os.O_TRUNC
}
writeFile, err := os.OpenFile(filename, openFlags, 0644)
if err != nil {
return err
}
defer writeFile.Close()
out = writeFile
} else {
out = os.Stdout
}
var err error
if tmpFile != nil {
_, err = io.CopyBuffer(out, tmpFile, buf)
} else {
_, err = out.Write(buf[:bufInitialRead])
}
return err
}
func cleanUp() {
if tmpFile == nil {
return
}
tmpFile.Close()
err := os.Remove(tmpFile.Name())
if err != nil {
log.Fatal(err)
}
}