//go:build linux
// +build linux

package main

import (
	"context"
	stderrors "errors"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"syscall"
	"time"

	"github.com/cenkalti/backoff/v4"
	cgroups1 "github.com/containerd/cgroups/v3/cgroup1"
	cgroups1stats "github.com/containerd/cgroups/v3/cgroup1/stats"
	oci "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/propagation"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"

	"github.com/Microsoft/hcsshim/internal/guest/bridge"
	"github.com/Microsoft/hcsshim/internal/guest/cgroup"
	"github.com/Microsoft/hcsshim/internal/guest/kmsg"
	"github.com/Microsoft/hcsshim/internal/guest/prot"
	"github.com/Microsoft/hcsshim/internal/guest/runtime/hcsv2"
	"github.com/Microsoft/hcsshim/internal/guest/runtime/runc"
	"github.com/Microsoft/hcsshim/internal/guest/transport"
	"github.com/Microsoft/hcsshim/internal/log"
	"github.com/Microsoft/hcsshim/internal/ot"
	"github.com/Microsoft/hcsshim/internal/version"
	"github.com/Microsoft/hcsshim/pkg/amdsevsnp"
	"github.com/Microsoft/hcsshim/pkg/securitypolicy"
)

func memoryLogFormat(metrics *cgroups1stats.Metrics) logrus.Fields {
	fields := logrus.Fields{}
	if metrics.Memory == nil {
		return fields
	}
	if metrics.Memory.Usage != nil {
		fields["memoryUsage"] = metrics.Memory.Usage.Usage
		fields["memoryUsageMax"] = metrics.Memory.Usage.Max
		fields["memoryUsageLimit"] = metrics.Memory.Usage.Limit
	}
	if metrics.Memory.Swap != nil {
		fields["swapUsage"] = metrics.Memory.Swap.Usage
		fields["swapUsageMax"] = metrics.Memory.Swap.Max
		fields["swapUsageLimit"] = metrics.Memory.Swap.Limit
	}
	if metrics.Memory.Kernel != nil {
		fields["kernelUsage"] = metrics.Memory.Kernel.Usage
		fields["kernelUsageMax"] = metrics.Memory.Kernel.Max
		fields["kernelUsageLimit"] = metrics.Memory.Kernel.Limit
	}
	return fields
}

func readMemoryEvents(startTime time.Time, efdFile *os.File, cgName string, threshold int64, mgr cgroup.Manager) {
	// Buffer must be >= 8 bytes for eventfd reads
	// http://man7.org/linux/man-pages/man2/eventfd.2.html
	count := 0
	buf := make([]byte, 8)
	isV1 := !cgroup.IsCgroupV2()
	for {
		if _, err := efdFile.Read(buf); err != nil {
			logrus.WithError(err).WithField("cgroup", cgName).Error("failed to read from eventfd")
			return
		}

		// For cgroup v1, check if event_control file still exists (for teardown detection).
		// cgroup v2 doesn't use event_control, so we skip this check.
		if isV1 {
			_, err := os.Lstat(fmt.Sprintf("/sys/fs/cgroup/memory%s/cgroup.event_control", cgName))
			if os.IsNotExist(err) {
				return
			}
		}

		count++
		msg := "memory usage for cgroup exceeded threshold"

		cgVersion := "v2"
		if isV1 {
			cgVersion = "v1"
		}
		entry := logrus.WithFields(logrus.Fields{
			"gcsStartTime":   startTime,
			"time":           time.Now(),
			"cgroup":         cgName,
			"thresholdBytes": threshold,
			"count":          count,
			"cgroup_version": cgVersion,
		})

		// Sleep for one second in case there is a series of allocations slightly after
		// reaching threshold.
		time.Sleep(time.Second)

		// Use the unified manager to get stats for both v1 and v2.
		metrics, err := mgr.Stats()
		if err != nil {
			// Don't return on Stat err as it will return an error if
			// any of the cgroup subsystems Stat calls failed for any reason.
			// We still want to log if we hit the cgroup threshold/limit
			entry.WithError(err).Error(msg)
		} else {
			entry.WithFields(memoryLogFormat(metrics)).Warn(msg)
		}
	}
}

// Sanity checks the memory limit value
func calculateContainersMemoryLimit(totalMemoryBytes, rootMemReserveBytes uint64) (int64, error) {
	if totalMemoryBytes <= rootMemReserveBytes {
		return 0, errors.Errorf("total memory %d must be greater than root memory reserve %d", totalMemoryBytes, rootMemReserveBytes)
	}

	limit := totalMemoryBytes - rootMemReserveBytes
	const maxInt64 = uint64(1<<63 - 1)
	if limit > maxInt64 {
		return 0, errors.Errorf("containers cgroup memory limit %d exceeds maximum supported value %d", limit, maxInt64)
	}
	return int64(limit), nil
}

// Returns the total system memory in bytes.
func currentTotalMemoryBytes() (uint64, error) {
	sinfo := syscall.Sysinfo_t{}
	if err := syscall.Sysinfo(&sinfo); err != nil {
		return 0, errors.Wrap(err, "failed to get sys info")
	}

	totalMemoryBytes := uint64(sinfo.Totalram)
	unit := uint64(sinfo.Unit)

	// If the unit is non-zero, multiply the total memory by the unit to get the actual memory in bytes.
	if unit != 0 {
		const maxUint64 = ^uint64(0)
		if totalMemoryBytes > maxUint64/unit {
			return 0, errors.Errorf("total memory %d with unit %d exceeds maximum supported value", totalMemoryBytes, unit)
		}
		totalMemoryBytes *= unit
	}
	return totalMemoryBytes, nil
}

// Start a periodic goroutine to monitor and update the pods cgroup memory limit.
func monitorCgroupMemoryLimits(h *hcsv2.Host, interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for range ticker.C {
		podLimit, updated, err := h.UpdateCgroupMemoryLimits()
		if err != nil {
			logrus.WithError(err).Error("failed to refresh cgroup memory limits")
			continue
		}
		if !updated {
			continue
		}
		logrus.WithField("memoryLimitBytes", podLimit).Debug("refreshed cgroup memory limits")
	}
}

// runWithRestartMonitor starts a command with given args and waits for it to exit. If the
// command exit code is non-zero the command is restarted with with some back off delay.
// Any stdout or stderr of the command will be split into lines and written as a log with
// logrus standard logger.  This function must be called in a separate goroutine.
func runWithRestartMonitor(arg0 string, args ...string) {
	backoffSettings := backoff.NewExponentialBackOff()
	// After we hit 10 min retry interval keep retrying after every 10 mins instead of
	// continuing to increase retry interval.
	backoffSettings.MaxInterval = time.Minute * 10
	for {
		command := exec.Command(arg0, args...)
		if err := command.Run(); err != nil {
			logrus.WithFields(logrus.Fields{
				"error":   err,
				"command": command.Args,
			}).Warn("restart monitor: run command returns error")
		}
		backOffTime := backoffSettings.NextBackOff()
		// since backoffSettings.MaxElapsedTime is set to 0 we will never receive backoff.Stop.
		time.Sleep(backOffTime)
	}
}

// startTimeSyncService starts the `chronyd` deamon to keep the UVM time synchronized.  We
// use a PTP device provided by the hypervisor as a source of correct time (instead of
// using a network server). We need to create a configuration file that configures chronyd
// to use the PTP device.  The system can have multiple PTP devices so we identify the
// correct PTP device by verifying that the `clock_name` of that device is `hyperv`.
func startTimeSyncService() error {
	ptpClassDir, err := os.Open("/sys/class/ptp")
	if err != nil {
		return errors.Wrap(err, "failed to open PTP class directory")
	}

	ptpDirList, err := ptpClassDir.Readdirnames(-1)
	if err != nil {
		return errors.Wrap(err, "failed to list PTP class directory")
	}

	var ptpDirPath string
	found := false
	// The file ends with a new line
	expectedClockName := "hyperv\n"
	for _, ptpDirPath = range ptpDirList {
		clockNameFilePath := filepath.Join(ptpClassDir.Name(), ptpDirPath, "clock_name")
		buf, err := os.ReadFile(clockNameFilePath)
		if err != nil && !os.IsNotExist(err) {
			return errors.Wrapf(err, "failed to read clock name file at %s", clockNameFilePath)
		}

		if string(buf) == expectedClockName {
			found = true
			break
		}
	}

	if !found {
		return errors.Errorf("no PTP device found with name \"%s\"", expectedClockName)
	}

	// create chronyd config file
	ptpDevPath := filepath.Join("/dev", filepath.Base(ptpDirPath))
	// chronyd config file take from: https://docs.microsoft.com/en-us/azure/virtual-machines/linux/time-sync
	chronydConfigString := fmt.Sprintf("refclock PHC %s poll 3 dpoll -2 offset 0 stratum 2\nmakestep 0.1 -1\n", ptpDevPath)
	chronydConfPath := "/tmp/chronyd.conf"
	err = os.WriteFile(chronydConfPath, []byte(chronydConfigString), 0644)
	if err != nil {
		return errors.Wrapf(err, "failed to create chronyd conf file %s", chronydConfPath)
	}

	// start chronyd. Do NOT start chronyd as daemon because creating a daemon
	// involves double forking the restart monitor will attempt to restart chornyd
	// after the first fork child exits.
	go runWithRestartMonitor("chronyd", "-n", "-f", chronydConfPath)
	return nil
}

func initOtelTracer() (func(context.Context) error, error) {
	exporter := &ot.LogrusExporter{}
	traceProvider := sdktrace.NewTracerProvider(
		sdktrace.WithSyncer(exporter),
		sdktrace.WithSampler(sdktrace.AlwaysSample()),
	)
	if traceProvider == nil {
		return nil, stderrors.New("failed to construct OpenTelemetry tracer provider")
	}
	otel.SetTracerProvider(traceProvider)
	if otel.GetTracerProvider() != traceProvider {
		return nil, stderrors.New("failed to register OpenTelemetry tracer provider globally")
	}
	otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(
		propagation.TraceContext{},
		propagation.Baggage{},
	))
	return traceProvider.Shutdown, nil
}

func main() {
	startTime := time.Now()
	logLevel := flag.String("loglevel",
		"debug",
		"Logging Level: debug, info, warning, error, fatal, panic.")
	coreDumpLoc := flag.String("core-dump-location",
		"",
		"The location/format where process core dumps will be written to.")
	kmsgLogLevel := flag.Uint("kmsgLogLevel",
		uint(kmsg.Warning),
		"Log all kmsg entries with a priority less than or equal to the supplied level.")
	logFile := flag.String("logfile",
		"",
		"Logging Target: An optional file name/path. Omit for console output.")
	logFormat := flag.String("log-format", "text", "Logging Format: text or json")
	useInOutErr := flag.Bool("use-inouterr",
		false,
		"If true use stdin/stdout for bridge communication and stderr for logging")
	v4 := flag.Bool("v4", false, "enable the v4 protocol support and v2 schema")
	rootMemReserveBytes := flag.Uint64("root-mem-reserve-bytes",
		75*1024*1024, // 75Mib
		"the amount of memory reserved for the orchestration, the rest will be assigned to containers")
	gcsMemLimitBytes := flag.Uint64("gcs-mem-limit-bytes",
		50*1024*1024, // 50 MiB
		"the maximum amount of memory the gcs can use")
	disableTimeSync := flag.Bool("disable-time-sync",
		false,
		"If true do not run chronyd time synchronization service inside the UVM")
	scrubLogs := flag.Bool("scrub-logs", true, "If true, scrub potentially sensitive information from logging")
	initialPolicyStance := flag.String("initial-policy-stance",
		"allow",
		"Stance: allow, deny.")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "\nUsage of %s:\n", os.Args[0])
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, "Examples:\n")
		fmt.Fprintf(os.Stderr, "    %s -loglevel=debug -logfile=/run/gcs/gcs.log\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "    %s -loglevel=info -logfile=stdout\n", os.Args[0])
	}

	flag.Parse()

	// If v4 enable otel tracing
	if *v4 {
		// Register our Otel logrus exporter
		shutdownTracer, err := initOtelTracer()
		if err != nil {
			logrus.Fatalf("failed to initialize ot tracer: %v", err)
		}
		defer func() { _ = shutdownTracer(context.Background()) }()
	}

	logrus.AddHook(log.NewHook())

	var logWriter *os.File
	if *logFile != "" {
		logFileHandle, err := os.OpenFile(*logFile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
		if err != nil {
			logrus.WithFields(logrus.Fields{
				"path":          *logFile,
				logrus.ErrorKey: err,
			}).Fatal("failed to create log file")
		}
		logWriter = logFileHandle
	} else {
		// logrus uses os.Stderr. see logrus.New()
		logWriter = os.Stderr
	}

	// set up our initial stance policy enforcer
	var initialEnforcer securitypolicy.SecurityPolicyEnforcer
	switch *initialPolicyStance {
	case "allow":
		initialEnforcer = &securitypolicy.OpenDoorSecurityPolicyEnforcer{}
		logrus.SetOutput(logWriter)
	case "deny":
		initialEnforcer = &securitypolicy.ClosedDoorSecurityPolicyEnforcer{}
		logrus.SetOutput(io.Discard)
	default:
		logrus.WithFields(logrus.Fields{
			"initial-policy-stance": *initialPolicyStance,
		}).Fatal("unknown initial-policy-stance")
	}

	switch *logFormat {
	case "text":
		// retain logrus's default.
	case "json":
		logrus.SetFormatter(&logrus.JSONFormatter{
			TimestampFormat: time.RFC3339Nano, // include ns for accurate comparisons on the host
		})
	default:
		logrus.WithFields(logrus.Fields{
			"log-format": *logFormat,
		}).Fatal("unknown log-format")
	}

	level, err := logrus.ParseLevel(*logLevel)
	if err != nil {
		logrus.Fatal(err)
	}

	logrus.SetLevel(level)

	log.SetScrubbing(*scrubLogs)

	logrus.WithFields(logrus.Fields{
		"branch":  version.Branch,
		"commit":  version.Commit,
		"version": version.Version,
	}).Info("GCS started")

	// Log which cgroup version is detected and will be used
	if cgroup.IsCgroupV2() {
		logrus.Info("cgroup v2 detected by GCS - using v2 API")
	} else {
		logrus.Info("cgroup v1 detected by GCS - using v1 API")
	}

	// Set the process core dump location. This will be global to all containers as it's a kernel configuration.
	// If no path is specified core dumps will just be placed in the working directory of wherever the process
	// was invoked to a file named "core".
	if *coreDumpLoc != "" {
		if err := os.WriteFile(
			"/proc/sys/kernel/core_pattern",
			[]byte(*coreDumpLoc),
			0644,
		); err != nil {
			logrus.WithError(err).Fatal("failed to set core dump location")
		}
	}

	// Continuously log /dev/kmsg
	go kmsg.ReadForever(kmsg.LogLevel(*kmsgLogLevel))

	// Setup the UVM cgroups to protect against a workload taking all available
	// memory and causing the GCS to malfunction we create cgroups: gcs,
	// containers, and pods for pod support.
	//

	// Write 1 to memory.use_hierarchy on the root cgroup to enable hierarchy
	// support. This needs to be set before we create any cgroups as the write
	// will fail otherwise. This is only needed for cgroup v1.
	//
	// On kernels where CONFIG_MEMCG_V1 is disabled (e.g. 6.18+), the v1
	// memory controller does not exist and this path will be absent even
	// though init fell back to v1 for other controllers. Without the memory
	// controller, container memory limits cannot be enforced and a workload
	// could OOM the entire UVM, so we still crash but with a diagnostic
	// message pointing at the kernel config.
	if !cgroup.IsCgroupV2() {
		if err := os.WriteFile("/sys/fs/cgroup/memory/memory.use_hierarchy", []byte("1"), 0644); err != nil {
			if os.IsNotExist(err) {
				logrus.WithError(err).Fatal("v1 memory controller not available; the kernel may lack CONFIG_MEMCG_V1 -- enable it or use a kernel with cgroup v2 + CONFIG_CGROUP_BPF")
			} else {
				logrus.WithError(err).Fatal("failed to enable hierarchy support for root cgroup")
			}
		}
	}

	// The pods cgroup is limited only by {Totalram - 75 MB (reservation)}.
	//
	// The gcs cgroup is not limited but an event will get logged if memory
	// usage exceeds 50 MB.
	totalMemoryBytes, err := currentTotalMemoryBytes()
	if err != nil {
		logrus.WithError(err).Fatal("failed to get total system memory")
	}
	podsLimit, err := calculateContainersMemoryLimit(totalMemoryBytes, *rootMemReserveBytes)
	if err != nil {
		logrus.WithError(err).Fatal("failed to calculate pods cgroup memory limit")
	}
	podsControl, err := cgroup.NewManager("/pods", &oci.LinuxResources{
		Memory: &oci.LinuxMemory{
			Limit: &podsLimit,
		},
	})
	if err != nil {
		logrus.WithError(err).Fatal("failed to create pods cgroup")
	}
	defer podsControl.Delete() //nolint:errcheck

	gcsControl, err := cgroup.NewManager("/gcs", &oci.LinuxResources{})
	if err != nil {
		logrus.WithError(err).Fatal("failed to create gcs cgroup")
	}
	defer gcsControl.Delete() //nolint:errcheck
	if err := gcsControl.Add(cgroups1.Process{Pid: os.Getpid()}); err != nil {
		logrus.WithError(err).Fatal("failed add gcs pid to gcs cgroup")
	}

	tport := &transport.VsockTransport{}
	rtime, err := runc.NewRuntime()
	if err != nil {
		logrus.WithError(err).Fatal("failed to initialize new runc runtime")
	}

	h := hcsv2.NewHost(rtime, tport, initialEnforcer, logWriter)
	h.InitializeCgroupMemoryLimitUpdater(podsControl, *rootMemReserveBytes, podsLimit)

	// Start periodic task to update the pods cgroup memory limit.
	go monitorCgroupMemoryLimits(h, 60*time.Second)

	// During live migration the VM is frozen and only wakes up when the host
	// shim is ready, so the vsock port should be immediately available. We
	// use a tight retry interval instead of exponential backoff.
	const reconnectInterval = 100 * time.Millisecond

	event := cgroups1.MemoryThresholdEvent(*gcsMemLimitBytes, false)
	gefd, err := gcsControl.RegisterMemoryEvent(event)
	if err != nil {
		logrus.WithError(err).Fatal("failed to register memory threshold for gcs cgroup")
	}
	gefdFile := os.NewFile(gefd, "gefd")
	defer gefdFile.Close()

	// Setup OOM monitoring for pods cgroup
	podsOom, err := podsControl.OOMEventFD()
	if err != nil {
		logrus.WithError(err).Fatal("failed to retrieve the pods cgroups oom eventfd")
	}
	podsOomFile := os.NewFile(podsOom, "pods-oomfd")
	defer podsOomFile.Close()

	// time synchronization service
	if !(*disableTimeSync) {
		if err = startTimeSyncService(); err != nil {
			logrus.WithError(err).Fatal("failed to start time synchronization service")
		}
	}

	go readMemoryEvents(startTime, gefdFile, "/gcs", int64(*gcsMemLimitBytes), gcsControl)
	go readMemoryEvents(startTime, podsOomFile, "/pods", podsLimit, podsControl)

	mux := bridge.NewBridgeMux()
	b := bridge.New(mux, *v4)
	// For confidential containers, we protect ourselves against attacks caused
	// by concurrent modifications, by processing one request at a time.
	sequential, err := amdsevsnp.IsSNP()
	if err != nil {
		// IsSNP cannot fail on LCOW
		log.G(context.TODO()).WithError(err).Error("unexpected error from IsSNP")
		// If it fails, we proceed with forceSequential enabled to be safe
		sequential = true
	}
	b.Sequential = sequential
	b.AssignHandlers(mux, h)

	// Reconnect loop: dial the host, serve until the connection drops, then
	// re-dial. During live migration the VM is frozen and only wakes up when
	// the destination host shim is ready, so the vsock port should be
	// immediately available.
	for {
		var bridgeIn io.ReadCloser
		var bridgeOut io.WriteCloser
		if *useInOutErr {
			bridgeIn = os.Stdin
			bridgeOut = os.Stdout
		} else {
			bridgeCon, dialErr := tport.Dial(prot.LinuxGcsVsockPort)
			if dialErr != nil {
				logrus.WithError(dialErr).Warn("failed to dial host, retrying")
				time.Sleep(reconnectInterval)
				continue
			}
			bridgeIn = bridgeCon
			bridgeOut = bridgeCon
		}

		logrus.Info("bridge connected, serving")

		serveErr := b.ListenAndServe(bridgeIn, bridgeOut)

		if b.ShutdownRequested() {
			logrus.Info("bridge shutdown requested, exiting reconnect loop")
			break
		}

		logrus.WithError(serveErr).Warn("bridge connection lost, will reconnect")
		time.Sleep(reconnectInterval)
	}
}