sponge/main.go

159 lines
3.0 KiB
Go
Raw Permalink Normal View History

2019-04-24 08:54:53 +00:00
// sponge - Soaks up all input from stdin and writes it to a file or stdout
2020-01-23 15:37:37 +00:00
// https://gitlab.com/tslocum/sponge
2019-04-24 08:54:53 +00:00
//
// 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"
2019-06-18 16:54:21 +00:00
"log"
2019-04-24 08:54:53 +00:00
"os"
2019-06-18 16:54:21 +00:00
"os/exec"
2019-04-24 08:54:53 +00:00
)
2019-06-22 12:44:41 +00:00
// 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.
2019-04-24 08:54:53 +00:00
const BufferSize = 64 * 1024 // 64 KiB
2019-06-22 12:44:41 +00:00
var (
buf = make([]byte, BufferSize)
bufInitialRead int
appendFile bool
tmpFile *os.File
)
2019-04-24 08:54:53 +00:00
2019-06-18 16:54:21 +00:00
func init() {
log.SetFlags(0)
log.SetPrefix("sponge: ")
}
2019-04-24 08:54:53 +00:00
2019-06-18 16:54:21 +00:00
func main() {
2019-06-22 12:44:41 +00:00
flag.BoolVar(&appendFile, "a", false, "append")
2019-04-24 08:54:53 +00:00
flag.Usage = func() {
2019-06-22 12:44:41 +00:00
fmt.Fprintf(os.Stderr, "%s [-a] [file]: soak up all input from stdin and write it to [file] or stdout\n", os.Args[0])
2019-04-24 08:54:53 +00:00
}
flag.Parse()
2019-06-22 12:44:41 +00:00
err := soak()
if err != nil {
log.Fatal(err)
}
defer cleanUp()
err = squeeze(flag.Arg(0))
if err != nil {
log.Fatal(err)
}
}
2019-04-24 08:54:53 +00:00
2019-06-22 12:44:41 +00:00
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
2019-06-18 16:54:21 +00:00
}
2019-04-24 08:54:53 +00:00
2019-06-22 12:44:41 +00:00
_, err = tmpFile.Write(buf[:bufInitialRead])
if err != nil {
tmpFile.Close()
return err
}
2019-04-24 08:54:53 +00:00
2019-06-22 12:44:41 +00:00
_, err = io.CopyBuffer(tmpFile, os.Stdin, buf)
if err != nil {
tmpFile.Close()
return err
}
2019-04-24 08:54:53 +00:00
2019-06-22 12:44:41 +00:00
_, err = tmpFile.Seek(io.SeekStart, 0)
if err != nil {
tmpFile.Close()
return err
}
}
2019-04-24 08:54:53 +00:00
2019-06-22 12:44:41 +00:00
return nil
}
2019-04-24 08:54:53 +00:00
2019-06-22 12:44:41 +00:00
func squeeze(filename string) error {
2019-06-18 16:54:21 +00:00
var out io.Writer
2019-06-22 12:44:41 +00:00
if filename != "" {
2019-06-18 16:54:21 +00:00
regularFile := true
2019-06-22 12:44:41 +00:00
fileInfo, err := os.Stat(filename)
2019-06-18 16:54:21 +00:00
if err != nil {
if !os.IsNotExist(err) {
2019-06-22 12:44:41 +00:00
return err
2019-06-18 16:54:21 +00:00
}
} else {
regularFile = fileInfo.Mode().IsRegular()
}
2019-06-22 12:44:41 +00:00
if tmpFile != nil && regularFile && !appendFile {
2019-06-18 16:54:21 +00:00
tmpFile.Close()
2019-06-22 12:44:41 +00:00
// Rename temporary file
err = os.Rename(tmpFile.Name(), filename)
2019-06-18 16:54:21 +00:00
if err != nil {
// Fall back to mv
2019-06-22 12:44:41 +00:00
cmd := exec.Command("mv", tmpFile.Name(), filename)
2019-06-18 16:54:21 +00:00
err = cmd.Run()
}
2019-06-22 12:44:41 +00:00
tmpFile = nil
return err
2019-06-18 16:54:21 +00:00
}
openFlags := os.O_WRONLY
if regularFile {
openFlags |= os.O_CREATE
}
2019-06-22 12:44:41 +00:00
if appendFile {
2019-06-18 16:54:21 +00:00
openFlags |= os.O_APPEND
2019-06-22 12:44:41 +00:00
} else if regularFile {
openFlags |= os.O_TRUNC
2019-06-18 16:54:21 +00:00
}
2019-06-22 12:44:41 +00:00
writeFile, err := os.OpenFile(filename, openFlags, 0644)
if err != nil {
return err
}
2019-06-18 16:54:21 +00:00
defer writeFile.Close()
2019-04-24 08:54:53 +00:00
2019-06-18 16:54:21 +00:00
out = writeFile
} else {
out = os.Stdout
2019-04-24 08:54:53 +00:00
}
2019-06-22 12:44:41 +00:00
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)
}
2019-04-24 08:54:53 +00:00
}