// -*- Mode: Go; indent-tabs-mode: t -*-

/*
 * Copyright (C) 2014-2020 Canonical Ltd
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 3 as
 * published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package osutil

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"syscall"
	"time"

	"github.com/snapcore/snapd/osutil/sys"
	"github.com/snapcore/snapd/randutil"
)

// AtomicWriteFlags are a bitfield of flags for AtomicWriteFile
type AtomicWriteFlags uint

const (
	// AtomicWriteFollow makes AtomicWriteFile follow symlinks
	AtomicWriteFollow AtomicWriteFlags = 1 << iota
)

// Allow disabling sync for testing. This brings massive improvements on
// certain filesystems (like btrfs) and very much noticeable improvements in
// all unit tests in genreal.
var snapdUnsafeIO bool = IsTestBinary() && GetenvBool("SNAPD_UNSAFE_IO", true)

// An AtomicFile is similar to an os.File but it has an additional
// Commit() method that does whatever needs to be done so the
// modification is "atomic": an AtomicFile will do its best to leave
// either the previous content or the new content in permanent
// storage. It also has a Cancel() method to abort and clean up.
type AtomicFile struct {
	*os.File

	target  string
	tmpname string
	uid     sys.UserID
	gid     sys.GroupID
	mtime   time.Time
	closed  bool
	renamed bool
}

// NewAtomicFile builds an AtomicFile backed by an *os.File that will have
// the given filename, permissions and uid/gid when Committed.
//
//	It _might_ be implemented using O_TMPFILE (see open(2)).
//
// Note that it won't follow symlinks and will replace existing symlinks with
// the real file, unless the AtomicWriteFollow flag is specified.
//
// It is the caller's responsibility to clean up on error, by calling Cancel().
//
// It is also the caller's responsibility to coordinate access to this, if it
// is used from different goroutines.
//
// Also note that there are a number of scenarios where Commit fails and then
// Cancel also fails. In all these scenarios your filesystem was probably in a
// rather poor state. Good luck.
func NewAtomicFile(filename string, perm os.FileMode, flags AtomicWriteFlags, uid sys.UserID, gid sys.GroupID) (aw *AtomicFile, err error) {
	if flags&AtomicWriteFollow != 0 {
		if fn, err := os.Readlink(filename); err == nil || (fn != "" && os.IsNotExist(err)) {
			if filepath.IsAbs(fn) {
				filename = fn
			} else {
				filename = filepath.Join(filepath.Dir(filename), fn)
			}
		}
	}
	// The tilde is appended so that programs that inspect all files in some
	// directory are more likely to ignore this file as an editor backup file.
	//
	// This fixes an issue in apparmor-utils package, specifically in
	// aa-enforce. Tools from this package enumerate all profiles by loading
	// parsing any file found in /etc/apparmor.d/, skipping only very specific
	// suffixes, such as the one we selected below.
	tmp := filename + "." + randutil.RandomString(12) + "~"

	fd, err := os.OpenFile(tmp, os.O_WRONLY|os.O_CREATE|os.O_TRUNC|os.O_EXCL, perm)
	if err != nil {
		return nil, err
	}

	return &AtomicFile{
		File:    fd,
		target:  filename,
		tmpname: tmp,
		uid:     uid,
		gid:     gid,
	}, nil
}

// ErrCannotCancel means the Commit operation failed at the last step, and
// your luck has run out.
var ErrCannotCancel = errors.New("cannot cancel: file has already been renamed")

func (aw *AtomicFile) Close() error {
	aw.closed = true
	return aw.File.Close()
}

// Cancel closes the AtomicWriter, and cleans up any artifacts. Cancel
// can fail if Commit() was (even partially) successful, but calling
// Cancel after a successful Commit does nothing beyond returning
// error--so it's always safe to defer a Cancel().
func (aw *AtomicFile) Cancel() error {
	if aw.renamed {
		return ErrCannotCancel
	}

	var e1, e2 error
	if aw.tmpname != "" {
		e1 = os.Remove(aw.tmpname)
	}
	if !aw.closed {
		e2 = aw.Close()
	}
	if e1 != nil {
		return e1
	}
	return e2
}

var chown = sys.Chown

const NoChown = sys.FlagID

// SetModTime sets the given modification time on the created file.
func (aw *AtomicFile) SetModTime(t time.Time) {
	aw.mtime = t
}

func (aw *AtomicFile) commit() error {
	if aw.uid != NoChown || aw.gid != NoChown {
		if err := chown(aw.File, aw.uid, aw.gid); err != nil {
			return err
		}
	}

	var dir *os.File
	if !snapdUnsafeIO {
		// XXX: if go switches to use aio_fsync, we need to open the dir for writing
		d, err := os.Open(filepath.Dir(aw.target))
		if err != nil {
			return err
		}
		dir = d
		defer dir.Close()

		if err := aw.Sync(); err != nil {
			return err
		}
	}

	if err := aw.Close(); err != nil {
		return err
	}

	if !aw.mtime.IsZero() {
		if err := os.Chtimes(aw.tmpname, time.Now(), aw.mtime); err != nil {
			return err
		}
	}

	if err := os.Rename(aw.tmpname, aw.target); err != nil {
		return err
	}
	aw.renamed = true // it is now too late to Cancel()

	if !snapdUnsafeIO {
		return dir.Sync()
	}

	return nil
}

// Commit the modification; make it permanent.
//
// If Commit succeeds, the writer is closed and further attempts to
// write will fail. If Commit fails, the writer _might_ be closed;
// Cancel() needs to be called to clean up.
func (aw *AtomicFile) Commit() error {
	return aw.commit()
}

// CommitAs commits the file under a new target name, following the same rules
// as Commit. The new target name must be located in the same directory as the
// original filename provided when creating AtomicFile.
//
// The call is useful when the target name is not known until the end (eg. it
// may depend on data being written to the file), in which case one can create
// AtomicFile using a temporary name and later override the actual name by
// calling CommitAs.
func (aw *AtomicFile) CommitAs(filename string) error {
	if dir := filepath.Dir(filename); dir != filepath.Dir(aw.target) {
		return fmt.Errorf("cannot commit as %q to a different directory %q", filepath.Base(filename), dir)
	}
	aw.target = filename
	return aw.commit()
}

// The AtomicWrite* family of functions work like os.WriteFile(), but the
// file created is an AtomicWriter, which is Committed before returning.
//
// AtomicWriteChown and AtomicWriteFileChown take an uid and a gid that can be
// used to specify the ownership of the created file. A special value of
// 0xffffffff (math.MaxUint32, or NoChown for convenience) can be used to
// request no change to that attribute.
//
// AtomicWriteFile and AtomicWriteFileChown take the content to be written as a
// []byte, and so work exactly like io.WriteFile(); AtomicWrite and
// AtomicWriteChown take an io.Reader which is copied into the file instead,
// and so are more amenable to streaming.
func AtomicWrite(filename string, reader io.Reader, perm os.FileMode, flags AtomicWriteFlags) (err error) {
	return AtomicWriteChown(filename, reader, perm, flags, NoChown, NoChown)
}

func AtomicWriteFile(filename string, data []byte, perm os.FileMode, flags AtomicWriteFlags) (err error) {
	return AtomicWriteChown(filename, bytes.NewReader(data), perm, flags, NoChown, NoChown)
}

func AtomicWriteFileChown(filename string, data []byte, perm os.FileMode, flags AtomicWriteFlags, uid sys.UserID, gid sys.GroupID) (err error) {
	return AtomicWriteChown(filename, bytes.NewReader(data), perm, flags, uid, gid)
}

func AtomicWriteChown(filename string, reader io.Reader, perm os.FileMode, flags AtomicWriteFlags, uid sys.UserID, gid sys.GroupID) (err error) {
	aw, err := NewAtomicFile(filename, perm, flags, uid, gid)
	if err != nil {
		return err
	}

	// Cancel once Committed is a NOP :-)
	defer aw.Cancel()

	if _, err := io.Copy(aw, reader); err != nil {
		return err
	}

	return aw.Commit()
}

// AtomicRename attempts to rename a path from oldName to newName atomically.
func AtomicRename(oldName, newName string) (err error) {
	var oldDir, newDir *os.File

	// snapdUnsafeIO controls the ability to ignore expensive disk
	// synchronization. It is only used inside tests.
	if !snapdUnsafeIO {
		// if called with a path with trailing '/', filepath.Dir
		// returns the dir itself instead of the parent
		oldName = filepath.Clean(oldName)
		newName = filepath.Clean(newName)

		oldDirPath := filepath.Dir(oldName)
		newDirPath := filepath.Dir(newName)

		oldDir, err = os.Open(oldDirPath)
		if err != nil {
			return err
		}
		defer oldDir.Close()

		newDir, err = os.Open(newDirPath)
		if err != nil {
			return err
		}
		defer newDir.Close()

		oldInfo, err := oldDir.Stat()
		if err != nil {
			return err
		}
		newInfo, err := newDir.Stat()
		if err != nil {
			return err
		}
		if oldStat, ok := oldInfo.Sys().(*syscall.Stat_t); ok {
			if newStat, ok := newInfo.Sys().(*syscall.Stat_t); ok {
				// Old and new directories refer to the same location. We can only sync once.
				if oldStat.Dev == newStat.Dev && oldStat.Ino == newStat.Ino {
					newDir = nil
				}
			}
		}
	}

	if err := os.Rename(oldName, newName); err != nil {
		return err
	}
	var err1, err2 error
	if oldDir != nil {
		err1 = oldDir.Sync()
	}
	if newDir != nil {
		err2 = newDir.Sync()
	}
	if err1 != nil {
		return err1
	}
	return err2
}

const maxLinkTries = 10

func atomicLinkOp(target, linkPath string, op func(target, tmp string) error, kind string) error {
	for tries := 0; tries < maxLinkTries; tries++ {
		tmp := linkPath + "." + randutil.RandomString(12) + "~"
		if err := op(target, tmp); err != nil {
			if os.IsExist(err) {
				continue
			}
			return err
		}
		defer os.Remove(tmp)
		return AtomicRename(tmp, linkPath)
	}
	return fmt.Errorf("cannot create a temporary %s", kind)
}

// AtomicSymlink attempts to atomically create a symlink at linkPath, pointing
// to a given target. The process creates a temporary symlink object pointing to
// the target, and then proceeds to rename it atomically, replacing the
// linkPath.
func AtomicSymlink(target, linkPath string) error {
	return atomicLinkOp(target, linkPath, os.Symlink, "symlink")
}

// AtomicLink attempts to atomically create a hardlink at linkPath, referencing
// the given target. The process creates a temporary hardlink object pointing
// to the target, and then proceeds to rename it atomically, replacing the
// linkPath.
func AtomicLink(target, linkPath string) error {
	return atomicLinkOp(target, linkPath, os.Link, "link")
}