package uffd

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"os"
	"sync/atomic"
	"syscall"
	"time"
	"unsafe"

	"github.com/buildbuddy-io/buildbuddy/enterprise/server/remote_execution/copy_on_write"
	"github.com/buildbuddy-io/buildbuddy/enterprise/server/remote_execution/snaputil"
	"github.com/buildbuddy-io/buildbuddy/server/metrics"
	"github.com/buildbuddy-io/buildbuddy/server/util/log"
	"github.com/buildbuddy-io/buildbuddy/server/util/status"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/sys/unix"
)

/*
#include <linux/userfaultfd.h> // For UFFD_API, UFFDIO_API
#include <linux/poll.h> // For POLLIN
*/
import "C"

// UFFD macros - see README for more info
const UFFDIO_COPY = 0xc028aa03
const UFFDIO_ZEROPAGE = 0xc020aa04

// uffdMsg is a notification from the userfaultfd object about a change in the
// virtual memory layout of the faulting process.
// It's defined in the linux header file userfaultfd.h
type UffdMsg = C.struct_uffd_msg

// A page fault event. Generated when the guest attempts to access a virtual
// memory page.
type Pagefault struct {
	Flags   uint64
	Address uint64
	Ptid    uint32
}

// Remove events are generated by Firecracker (via MADV_DONTNEED) when the
// balloon expands into memory, to mark it as unused.
type Remove struct {
	Start uint64
	End   uint64
}

type uffdEvent struct {
	pagefault *Pagefault
	remove    *Remove
}

// uffdioCopy contains input/output data to the UFFDIO_COPY syscall.
//
// It can be used to resolve a page fault by copying data from some source (such
// as a memory snapshot file) to a destination address (such as into the guest).
type uffdioCopy struct {
	Dst  uint64 // Address of the faulting region that memory should be copied to
	Src  uint64 // Address of the data to copy
	Len  uint64 // Number of bytes to copy
	Mode uint64 // Flags controlling behavior of copy
	Copy int64  // After the syscall has completed, contains the number of bytes copied, or a negative errno to indicate failure
}

// uffdioZeropage contains input and output data to the UFFDIO_ZEROPAGE syscall.
//
// It can be used to resolve a page fault when all 0s should be faulted in.
// Empty pages are expected by the guest after the balloon has removed them.
type uffdIoZeropage struct {
	Range uffdIoRange
	Mode  uint64
	// After the syscall has completed, contains the number of bytes zeroed, or a negative errno to indicate failure.
	Zeropage int64
}

type uffdIoRange struct {
	Start uint64
	Len   uint64
}

// GuestRegionUFFDMapping represents the mapping between a VM memory address to the offset in the corresponding
// memory snapshot file.
//
// # This data is sent by firecracker to tell the UFFD handler how to handle page faults
//
// Ex. If the VM tries to read memory at `BaseHostVirtAddr` and triggers a page fault, the UFFD handler can populate it
// by copying `size` bytes from `offset` in the memory snapshot file.
type GuestRegionUFFDMapping struct {
	BaseHostVirtAddr uintptr `json:"base_host_virt_addr"`
	Size             uintptr `json:"size"`
	Offset           uintptr `json:"offset"`
}

func (g *GuestRegionUFFDMapping) ContainsGuestAddr(addr uintptr) bool {
	return addr >= g.BaseHostVirtAddr && addr < g.BaseHostVirtAddr+g.Size
}

// Initial message sent on the socket by firecracker.
type setupMessage struct {
	Mappings []GuestRegionUFFDMapping
	Uffd     uintptr
}

// When loading a firecracker memory snapshot, this userfaultfd handler can be used to handle page faults for the VM.
// This handler uses a copy_on_write.COWStore to manage the snapshot - allowing it to be served remotely, compressed, etc.
type Handler struct {
	lis            net.Listener
	stoppedChan    chan struct{}
	handleDoneChan chan struct{}
	handleErr      error

	earlyTerminationReader *os.File
	earlyTerminationWriter *os.File

	pageFaultTotalDurationNanos atomic.Int64

	// Addresses of pages that were removed by the balloon.
	// If these pages page fault again, they should be handled with UFFDIO_ZEROPAGE.
	removedAddresses map[int64]struct{}
}

func NewHandler() (*Handler, error) {
	return &Handler{
		removedAddresses: make(map[int64]struct{}),
	}, nil
}

// Start starts a goroutine to listen on the given socket path for Firecracker's
// UFFD initialization message, and then starts fulfilling UFFD requests using
// the given memory store.
func (h *Handler) Start(ctx context.Context, socketPath string, memoryStore *copy_on_write.COWStore) error {
	lis, err := net.ListenUnix("unix", &net.UnixAddr{Name: socketPath, Net: "unix"})
	if err != nil {
		return status.WrapError(err, "listen on socket")
	}
	h.lis = lis
	log.CtxDebugf(ctx, "userfaultfd handler listening on unix://%s", socketPath)

	// Set the permissions of the socket file
	if err := os.Chmod(socketPath, 0777); err != nil {
		return status.WrapError(err, "set socket permissions")
	}

	// receiveSetupMsg() will terminate early if Stop() is called before
	// receiving the UFFD. This may happen if firecracker crashes due to failing
	// to initialize the UFFD object.
	h.stoppedChan = make(chan struct{})
	// Stop() will wait until this is closed to verify the handler has completed
	// handling open requests
	h.handleDoneChan = make(chan struct{})

	// Create a FD that can be used to terminate Poll early
	pipeRead, pipeWrite, err := os.Pipe()
	if err != nil {
		return status.WrapError(err, "create early-termination fd")
	}
	h.earlyTerminationReader = pipeRead
	h.earlyTerminationWriter = pipeWrite

	go func() {
		h.handleErr = h.handle(ctx, memoryStore)
		if h.handleErr != nil {
			log.CtxErrorf(ctx, "UFFD handler failed: %s", h.handleErr)
		}
		close(h.handleDoneChan)
	}()
	return nil
}

// When UFFD is set as the memory backend, firecracker will create a UFFD object and send it to
// the UFFD handler over a unix socket. It also sends GuestRegionUFFDMapping data so the handler knows
// how to resolve the page faults
func (h *Handler) receiveSetupMsg(ctx context.Context) (*setupMessage, error) {
	setupDone := make(chan struct{})
	defer close(setupDone)
	go func() {
		select {
		case <-h.stoppedChan:
		case <-setupDone:
		}
		h.lis.Close()
	}()

	log.CtxDebugf(ctx, "Waiting for firecracker to connect to uffd socket")
	conn, err := h.lis.Accept()
	if err != nil {
		select {
		case <-h.stoppedChan:
			return nil, status.InternalError("handler stopped while waiting to receive setup message")
		default:
			return nil, status.WrapError(err, "accept firecracker connection")
		}
	}
	unixConn := conn.(*net.UnixConn)
	log.CtxDebugf(ctx, "Firecracker connected to uffd socket")

	// Read data sent from firecracker.
	// The GuestRegionUFFDMapping data is serialized as JSON. The UFFD object it created is sent as "out-of-band" data.
	mappingsBuf := make([]byte, 1024)
	// The size of a FD is 4B - only 1 should be sent (the UFFD object)
	uffdBuf := make([]byte, syscall.CmsgSpace(4))

	numBytesMappings, numBytesFD, _, _, err := unixConn.ReadMsgUnix(mappingsBuf, uffdBuf)
	if err != nil {
		return nil, status.WrapError(err, "failed to read unix msg from connection")
	}

	// Parse memory mappings
	mappingsBuf = mappingsBuf[:numBytesMappings]
	var mappings []GuestRegionUFFDMapping
	err = json.Unmarshal(mappingsBuf, &mappings)
	if err != nil {
		return nil, status.WrapError(err, "parse memory mapping data")
	}
	log.CtxDebugf(ctx, "Received memory region mappings: %s", string(mappingsBuf))

	// Parse UFFD object
	controlMsgs, err := syscall.ParseSocketControlMessage(uffdBuf[:numBytesFD])
	if err != nil {
		return nil, status.WrapError(err, "parse control messages")
	}
	if len(controlMsgs) != 1 {
		return nil, status.InternalErrorf("expected 1 control message containing UFFD, found %d", len(controlMsgs))
	}
	fds, err := syscall.ParseUnixRights(&controlMsgs[0])
	if err != nil {
		return nil, status.WrapError(err, "parse unix writes")
	}
	if len(fds) != 1 {
		return nil, status.InternalErrorf("expected 1 fd (the uffd object), found %d", len(fds))
	}
	uffd := uintptr(fds[0])

	return &setupMessage{
		Uffd:     uffd,
		Mappings: mappings,
	}, nil
}

func (h *Handler) handle(ctx context.Context, memoryStore *copy_on_write.COWStore) (err error) {
	// Get uffd sent from firecracker
	setup, err := h.receiveSetupMsg(ctx)
	if err != nil {
		return status.WrapError(err, "receive setup message from firecracker")
	}
	uffd := setup.Uffd
	mappings := setup.Mappings
	defer syscall.Close(int(uffd))

	pollFDs := []unix.PollFd{
		{Fd: int32(uffd), Events: C.POLLIN},
		{Fd: int32(h.earlyTerminationReader.Fd()), Events: C.POLLIN},
	}
	pageSize := os.Getpagesize()

	deferredPageFaultEvents := make([]*Pagefault, 0)
	for {
		// Poll UFFD for messages
		_, pollErr := unix.Poll(pollFDs, -1)
		if pollErr != nil {
			if pollErr == unix.EINTR {
				// Poll call was interrupted by another signal - retry
				continue
			}
			return status.WrapError(pollErr, "poll uffd")
		}

		// Check for an early termination message
		earlyTerminationFd := pollFDs[1]
		if earlyTerminationFd.Revents&C.POLLIN != 0 {
			return nil
		}

		// Process page faults we couldn't handle last round.
		uffdEvents := make([]*uffdEvent, 0)
		for _, pf := range deferredPageFaultEvents {
			uffdEvents = append(uffdEvents, &uffdEvent{pagefault: pf})
		}
		deferredPageFaultEvents = make([]*Pagefault, 0)

		// Read all available UFFD notifications.
		//
		// We need to read all notifications because if there is an unread
		// remove event in the queue, all attempts to resolve page faults with
		// UFFDIO_COPY and UFFDIO_ZEROPAGE will fail with EAGAIN and the handler
		// will be blocked.
		for {
			event, err := readEvent(uffd)
			if err != nil {
				if errors.Is(err, unix.EAGAIN) {
					// No more messages in the notification queue
					break
				}
				return status.InternalErrorf("read event from uffd failed with errno(%d)", err)
			}

			uffdEvents = append(uffdEvents, event)
		}

		// First, handle all remove events.
		//
		// Page faults and remove events are not guaranteed to be ordered, because
		// page faults are triggered within the guest on the vCPU thread, while
		// remove events are triggered by Firecracker on the host. This can
		// create a race condition (for example if deflate_on_oom is set, the guest
		// may try to reclaim a page immediately after the balloon has removed it).
		//
		// If there is a race, we should always zero the page. If the balloon
		// removed the page first, the subsequent page fault should be zeroed.
		// Even if the page fault occurs first, the balloon immediately removes it,
		// so the page should still be zeroed.
		//
		// To guarantee that any remove events are handled before page faults on
		// the same address, handle all remove events first.
		pageFaults := make([]*Pagefault, 0)
		for _, e := range uffdEvents {
			if e.remove != nil {
				if valid := validateRemoveEvent(e.remove, pageSize); !valid {
					return status.InternalErrorf("unexpected remove event %v", e.remove)
				}

				affectedChunks := make(map[int64]struct{})
				for i := int64(e.remove.Start); i < int64(e.remove.End); i += int64(os.Getpagesize()) {
					h.removedAddresses[i] = struct{}{}

					// Track affected chunks
					storeOffset, err := storeOffsetForGuestAddr(uintptr(i), mappings)
					if err != nil {
						return err
					}
					chunkStartOffset := memoryStore.ChunkStartOffset(storeOffset)
					affectedChunks[chunkStartOffset] = struct{}{}
				}

				// Mark affected chunks as partially mapped
				for chunkStartOffset, _ := range affectedChunks {
					memoryStore.MarkPartiallyMapped(chunkStartOffset)
				}
			} else if e.pagefault != nil {
				pageFaults = append(pageFaults, e.pagefault)
			}
		}

		// Handle page faults.
		deferPageFaults := false
		for _, pf := range pageFaults {
			if deferPageFaults {
				deferredPageFaultEvents = append(deferredPageFaultEvents, pf)
			} else {
				// The memory location the VM tried to access that triggered the page fault.
				guestFaultingAddr := pf.Address

				if err := h.handlePageFault(uffd, memoryStore, mappings, uintptr(guestFaultingAddr)); err != nil {
					// If a new remove event has arrived in the queue before we finish
					// handling all the page faults, the remaining page faults will
					// all fail with EAGAIN. Defer these page faults so we can
					// handle them after the notification queue is empty again.
					if errors.Is(err, unix.EAGAIN) {
						deferPageFaults = true
						deferredPageFaultEvents = append(deferredPageFaultEvents, pf)
					} else {
						return err
					}
				}
			}
		}
	}
}

// EmitSummaryMetrics will export cumulative metrics for the given stage, and
// then reset them.
func (h *Handler) EmitSummaryMetrics(stage string) {
	// Read the total accumulated duration and reset the counter to 0.
	pageFaultTotalDuration := time.Duration(h.pageFaultTotalDurationNanos.Swap(0))
	metrics.COWSnapshotPageFaultTotalDurationUsec.With(prometheus.Labels{
		metrics.Stage: stage,
	}).Observe(float64(pageFaultTotalDuration.Microseconds()))
}

// handlePageFault prepares for and resolves a page fault.
func (h *Handler) handlePageFault(uffd uintptr, memoryStore *copy_on_write.COWStore, mappings []GuestRegionUFFDMapping, faultingAddress uintptr) error {
	mapping, err := guestMemoryAddrToMapping(faultingAddress, mappings)
	if err != nil {
		return err
	}

	// We can only handle page faults in units of pages, so align the faulting
	// address to a page boundary. In other words, get the address of the
	// start of the page containing the faulting address.
	guestPageAddr := pageStartAddress(uint64(faultingAddress), os.Getpagesize())
	if guestPageAddr < mapping.BaseHostVirtAddr {
		// Make sure we only try to map addresses that fall within the valid
		// guest memory ranges
		guestPageAddr = mapping.BaseHostVirtAddr
	}

	// If address had been previously removed, zero it.
	if _, removed := h.removedAddresses[int64(guestPageAddr)]; removed {
		return h.copyZeroes(uffd, guestPageAddr)
	}

	// Otherwise copy data from the memory snapshot.
	return h.copyDataFromMemorySnapshot(uffd, memoryStore, mapping, guestPageAddr)
}

// copyZeroes handles a page fault by copying zeroes into the guest. The guest
// expects zeroes if the faulting page was previously removed by the balloon.
func (h *Handler) copyZeroes(uffd uintptr, faultingAddress uintptr) error {
	// If multiple consecutive pages were removed, zero them all.
	end := faultingAddress
	for {
		_, removed := h.removedAddresses[int64(end)]
		if !removed {
			break
		}

		end += uintptr(os.Getpagesize())
	}

	sizeToZero := end - faultingAddress
	if sizeToZero > 0 {
		zeroIO := uffdIoZeropage{
			Range: uffdIoRange{
				Start: uint64(faultingAddress),
				Len:   uint64(sizeToZero),
			},
		}
		_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uffd, UFFDIO_ZEROPAGE, uintptr(unsafe.Pointer(&zeroIO)))
		if errno != 0 {
			return wrapErrno(errno, fmt.Sprintf("UFFDIO_ZEROPAGE failed with errno(%d)", errno))
		}
	}

	for zeroedPage := faultingAddress; zeroedPage < end; zeroedPage += uintptr(os.Getpagesize()) {
		delete(h.removedAddresses, int64(zeroedPage))
	}
	return nil
}

// copyDataFromMemorySnapshot handles a page fault by copying data from the memory
// snapshot into the guest.
func (h *Handler) copyDataFromMemorySnapshot(uffd uintptr, memoryStore *copy_on_write.COWStore, mapping *GuestRegionUFFDMapping, faultingAddress uintptr) error {
	// Find the memory data in the store that should be used to handle the page fault
	faultStoreOffset := guestMemoryAddrToStoreOffset(faultingAddress, *mapping)
	chunkStartOffset := memoryStore.ChunkStartOffset(int64(faultStoreOffset))

	var hostAddr, destAddr uintptr
	var copySize int64
	var err error
	if *snaputil.EnableBalloon && memoryStore.PartiallyMapped(chunkStartOffset) {
		// If the chunk is partially mapped (i.e. contains some removed
		// pages), only page fault in 1 page at a time
		// because holes and data must be handled differently and chunk-level
		// operations don't provide enough granularity.
		hostAddr, err = memoryStore.GetPageAddress(faultStoreOffset, false)
		if err != nil {
			return status.WrapError(err, "get backing page address")
		}
		destAddr = faultingAddress
		copySize = int64(os.Getpagesize())
	} else {
		//To reduce the number of UFFD round trips, try to copy the entire
		//chunk containing the faulting address.
		hostAddr, copySize, err = memoryStore.GetChunkStartAddressAndSize(faultStoreOffset, false /*=write*/)
		if err != nil {
			return status.WrapError(err, "get backing page address")
		}
		relOffset := memoryStore.GetRelativeOffsetFromChunkStart(faultStoreOffset)
		destAddr = faultingAddress - relOffset
	}

	// Check that the request is within the valid guest memory range.
	//
	// Check for data below the valid memory range.
	var invalidBytesAtChunkStart uintptr
	if destAddr < mapping.BaseHostVirtAddr {
		invalidBytesAtChunkStart = mapping.BaseHostVirtAddr - destAddr
		destAddr = mapping.BaseHostVirtAddr
		hostAddr += invalidBytesAtChunkStart
		copySize -= int64(invalidBytesAtChunkStart)
	}
	// Check for data above the valid memory range
	mappingEndAddr := mapping.BaseHostVirtAddr + mapping.Size
	copyEndAddr := destAddr + uintptr(copySize)
	var invalidBytesAtChunkEnd int64
	if copyEndAddr > mappingEndAddr {
		invalidBytesAtChunkEnd = int64(copyEndAddr - mappingEndAddr)
		copySize -= invalidBytesAtChunkEnd
	}

	_, err = h.resolvePageFault(uffd, uint64(destAddr), uint64(hostAddr), uint64(copySize))
	if err != nil {
		return err
	}

	// After memory has been copied to the VM, unmap the chunk to save memory
	// usage on the executor
	if err := memoryStore.UnmapChunk(chunkStartOffset); err != nil {
		return err
	}
	return nil
}

// resolvePageFault executes the resolution of a page fault.
// It copies `size` bytes of memory from a `Src` address to the faulting region `Dst`.
//
// When complete, the UFFDIO_COPY call wakes up the process that triggered the page fault (When a process
// attempts to access unallocated memory, it triggers a page fault and hangs until it has been resolved)
//
// Returns the number of bytes copied
func (h *Handler) resolvePageFault(uffd uintptr, faultingRegion uint64, src uint64, size uint64) (int64, error) {
	start := time.Now()
	defer func() {
		h.pageFaultTotalDurationNanos.Add(time.Since(start).Nanoseconds())
	}()

	copyData := uffdioCopy{
		Dst: faultingRegion,
		Src: src,
		Len: size,
	}
	_, _, errno := syscall.Syscall(syscall.SYS_IOCTL, uffd, UFFDIO_COPY, uintptr(unsafe.Pointer(&copyData)))
	if errno != 0 {
		// If error is due to the page already being mapped, ignore.
		if errno == unix.EEXIST {
			return 0, nil
		}
		return 0, wrapErrno(errno, fmt.Sprintf("UFFDIO_COPY failed with errno(%d)", errno))
	}
	return copyData.Copy, nil
}

// readEvent reads a notification from UFFD.
func readEvent(uffd uintptr) (*uffdEvent, error) {
	var event UffdMsg
	_, _, errno := syscall.Syscall(syscall.SYS_READ, uffd, uintptr(unsafe.Pointer(&event)), unsafe.Sizeof(event))
	if errno != 0 {
		return nil, errno
	}

	if event.event == C.UFFD_EVENT_PAGEFAULT {
		pagefault := (*(*Pagefault)(unsafe.Pointer(&event.arg[0])))
		return &uffdEvent{pagefault: &pagefault}, nil
	} else if event.event == C.UFFD_EVENT_REMOVE {
		remove := (*(*Remove)(unsafe.Pointer(&event.arg[0])))
		return &uffdEvent{remove: &remove}, nil
	}
	return nil, status.InternalErrorf("unsupported uffd event type %v", event.event)
}

// Gets the address of the start of the memory page containing `addr`
func pageStartAddress(addr uint64, pageSize int) uintptr {
	// Align the address to the nearest lower multiple of pageSize by masking the least significant bits.
	return uintptr(addr & ^(uint64(pageSize) - 1))
}

// storeOffsetForGuestAddr returns the address in the store that corresponds
// to the given guest memory address.
func storeOffsetForGuestAddr(guestAddr uintptr, memoryMappings []GuestRegionUFFDMapping) (int64, error) {
	mapping, err := guestMemoryAddrToMapping(guestAddr, memoryMappings)
	if err != nil {
		return 0, err
	}

	// Find the memory data in the store that should be used to handle the page fault
	storeOffset := guestMemoryAddrToStoreOffset(guestAddr, *mapping)
	return int64(storeOffset), nil
}

func validateRemoveEvent(remove *Remove, pageSize int) bool {
	pageAligned := remove.Start%uint64(pageSize) == 0 && remove.End%uint64(pageSize) == 0
	return pageAligned
}

// Wait waits for the UFFD handler to terminate. It returns an error only if
// the handler terminated because it failed to handle a page fault request.
func (h *Handler) Wait() error {
	<-h.handleDoneChan
	return h.handleErr
}

func (h *Handler) Stop() error {
	if h.earlyTerminationWriter == nil || h.earlyTerminationReader == nil || h.handleDoneChan == nil {
		log.Info("UFFD handler was already stopped when Stop() was called")
		return nil
	}

	close(h.stoppedChan)

	log.Info("UFFD handler beginning shut down")
	_, err := h.earlyTerminationWriter.Write([]byte{0})
	if err != nil {
		return status.WrapError(err, "write to early terminator")
	}

	// Wait for handle() to finish processing open requests
	<-h.handleDoneChan

	h.earlyTerminationReader.Close()
	h.earlyTerminationWriter.Close()
	h.earlyTerminationReader = nil
	h.earlyTerminationWriter = nil
	return nil
}

// Translate the faulting memory address in the guest to a persisted store offset
// based on the memory mappings.
func guestMemoryAddrToStoreOffset(addr uintptr, mapping GuestRegionUFFDMapping) uintptr {
	relativeOffset := addr - mapping.BaseHostVirtAddr
	return mapping.Offset + relativeOffset
}

// Return the GuestRegionUFFDMapping corresponding to the given address
func guestMemoryAddrToMapping(addr uintptr, mappings []GuestRegionUFFDMapping) (*GuestRegionUFFDMapping, error) {
	for _, m := range mappings {
		if !m.ContainsGuestAddr(addr) {
			continue
		}
		return &m, nil
	}
	return nil, status.InternalErrorf("page address 0x%x not found in guest region UFFD mappings", addr)
}

// wrapErrno wraps an errno with a message, while still preserving the
// underlying error type so that it can be compared with errors.Is().
//
// status.WrapError only extracts the original error's message as a string, but
// does not preserve its type.
func wrapErrno(err syscall.Errno, msg string) error {
	return fmt.Errorf("%s: %w", msg, err)
}
