// Copyright 2020 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package fuzz

import (
	"bytes"
	"os"
	"os/exec"
	"strings"
	"testing"
	"time"

	"github.com/google/go-cmp/cmp"
	"github.com/google/go-cmp/cmp/cmpopts"
)

const lateBootMessage = "we made it!"

func TestQemuLauncherHandle(t *testing.T) {
	launcher := &QemuLauncher{Pid: 4141, TmpDir: "/tmp/woof"}

	handle, err := NewHandleWithData(HandleData{launcher: launcher})
	if err != nil {
		t.Fatalf("error creating handle: %s", err)
	}
	defer handle.Release()

	// Note: we don't serialize here because that is covered by handle tests

	build, _ := newMockBuild()

	reloadedLauncher, err := loadLauncherFromHandle(build, handle, true)
	if err != nil {
		t.Fatalf("error loading launcher from handle: %s", err)
	}

	ql, ok := reloadedLauncher.(*QemuLauncher)
	if !ok {
		t.Fatalf("incorrect launcher type")
	}

	if diff := cmp.Diff(launcher, ql, cmpopts.IgnoreUnexported(QemuLauncher{})); diff != "" {
		t.Fatalf("incorrect data in reloaded launcher (-want +got):\n%s", diff)
	}
}

func TestIncompleteQemuLauncherHandle(t *testing.T) {
	// Construct an object that isn't fully initialized
	launcher := &QemuLauncher{Pid: 4141}

	handle, err := NewHandleWithData(HandleData{launcher: launcher})
	if err != nil {
		t.Fatalf("error creating handle: %s", err)
	}
	defer handle.Release()

	build, _ := newMockBuild()

	if _, err := loadConnectorFromHandle(build, handle, true); err == nil {
		t.Fatalf("expected error, but succeeded")
	}
}

func TestQemuLauncherWithMissingDeps(t *testing.T) {
	// Enable subprocess mocking
	ExecCommand = mockCommand
	defer func() { ExecCommand = exec.Command }()

	for _, dep := range []string{"zbi", "fvm", "blk", "zbitool", "qemu", "kernel"} {
		// This awkward error-discarding and casting is due to the need for the
		// signature of newMockBuild to match the other Build-creating functions
		build, _ := newMockBuild()
		brokenBuild := build.(*mockBuild)
		defer os.RemoveAll(brokenBuild.enableQemu(t, FakeQemuNormal))

		brokenBuild.paths[dep] = invalidPath
		launcher := NewQemuLauncher(build)
		if _, err := launcher.Start(); err == nil {
			t.Fatalf("Expected error launching with missing %q dep but succeeded", dep)
		}

		expectPathAbsent(t, launcher.TmpDir)
	}
}

func TestQemuLauncherWithQemuFailure(t *testing.T) {
	// Enable subprocess mocking
	ExecCommand = mockCommand
	defer func() { ExecCommand = exec.Command }()

	build, _ := newMockBuild()
	defer os.RemoveAll(build.(*mockBuild).enableQemu(t, FakeQemuFailing))

	launcher := NewQemuLauncher(build)
	if _, err := launcher.Start(); err == nil {
		t.Fatalf("Expected error launching with failing qemu but succeeded")
	}

	expectPathAbsent(t, launcher.TmpDir)
}

func TestQemuLauncherWithQemuTimeout(t *testing.T) {
	// Enable subprocess mocking
	ExecCommand = mockCommand
	defer func() { ExecCommand = exec.Command }()

	build, _ := newMockBuild()
	defer os.RemoveAll(build.(*mockBuild).enableQemu(t, FakeQemuSlow))

	launcher := NewQemuLauncher(build)
	launcher.timeout = 10 * time.Millisecond
	if _, err := launcher.Start(); err == nil {
		t.Fatalf("Expected error launching with slow qemu but succeeded")
	}

	if running, err := launcher.IsRunning(); err != nil || running {
		t.Fatalf("expected qemu to be killed after timeout")
	}

	expectPathAbsent(t, launcher.TmpDir)
}

// Helper for detecting changes in files generated by Prepare
func getFileHashes(t *testing.T, q *QemuLauncher) map[string][]byte {
	hashes := make(map[string][]byte)
	for _, path := range []string{q.initrd, q.extendedBlk, q.sshKey, q.sshPublicKey} {
		hashes[path] = getFileHash(t, path)
	}
	return hashes
}

func TestQemuLauncherPrepare(t *testing.T) {
	// Enable subprocess mocking
	ExecCommand = mockCommand
	defer func() { ExecCommand = exec.Command }()

	build, _ := newMockBuild()
	launcher := NewQemuLauncher(build)

	if err := launcher.Prepare(); err != nil {
		t.Fatalf("Error preparing launcher: %s", err)
	}

	defer launcher.Kill()

	origFileHashes := getFileHashes(t, launcher)

	// Nothing should change if re-run
	if err := launcher.Prepare(); err != nil {
		t.Fatalf("Error re-preparing launcher: %s", err)
	}

	fileHashes := getFileHashes(t, launcher)

	if diff := cmp.Diff(origFileHashes, fileHashes); diff != "" {
		t.Fatalf("Prepare was not idempotent (-want +got):\n%s", diff)
	}

	// Check for regeneration of SSH key when necessary
	if err := os.Remove(launcher.sshKey); err != nil {
		t.Fatalf("Error deleting ssh key: %s", err)
	}

	if err := launcher.Prepare(); err != nil {
		t.Fatalf("Error re-preparing launcher: %s", err)
	}

	// ZBI should be updated if SSH key is regenerated
	shouldChange := []string{launcher.sshKey, launcher.sshPublicKey, launcher.initrd}
	for _, path := range shouldChange {
		if bytes.Equal(getFileHash(t, path), origFileHashes[path]) {
			t.Fatalf("File not updated after SSH key regeneration: %q", path)
		}
	}
}

func TestQemuLauncher(t *testing.T) {
	// Enable subprocess mocking
	ExecCommand = mockCommand
	defer func() { ExecCommand = exec.Command }()

	build, _ := newMockBuild()
	defer os.RemoveAll(build.(*mockBuild).enableQemu(t, FakeQemuNormal))

	launcher := NewQemuLauncher(build)
	if _, err := launcher.Start(); err != nil {
		t.Fatalf("Error starting instance: %s", err)
	}

	alive, err := launcher.IsRunning()
	if err != nil {
		t.Fatalf("Error checking instance status: %s", err)
	}

	if !alive {
		t.Fatalf("instance not alive when it should be")
	}

	var out bytes.Buffer
	if err := launcher.GetLogs(&out); err != nil {
		t.Fatalf("Error getting instance logs: %s", err)
	}
	if !strings.Contains(out.String(), lateBootMessage) {
		t.Fatalf("Instance logs (%d bytes) missing late entry", out.Len())
	}

	if err := launcher.Kill(); err != nil {
		t.Fatalf("Error killing instance: %s", err)
	}

	alive, err = launcher.IsRunning()
	if err != nil {
		t.Fatalf("Error checking instance status: %s", err)
	}

	if alive {
		t.Fatalf("instance alive when it shouldn't be")
	}

	expectPathAbsent(t, launcher.TmpDir)
}