<commit_before>/*
Copyright 2020 gRPC authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"context"
	"fmt"
	"time"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("WaitForReadyPods", func() {
	var fastDuration time.Duration
	var slowDuration time.Duration

	var irrelevantPod corev1.Pod
	var driverPod corev1.Pod
	var serverPod corev1.Pod
	var clientPod corev1.Pod

	BeforeEach(func() {
		fastDuration = 1 * time.Millisecond * timeMultiplier
		slowDuration = 100 * time.Millisecond * timeMultiplier

		irrelevantPod = corev1.Pod{}

		driverPod = newTestPod("driver")
		driverPod.Spec.Containers[0].Ports = nil

		serverPod = newTestPod("server")
		serverPod.Status.PodIP = "127.0.0.2"

		clientPod = newTestPod("client")
		clientPod.Status.PodIP = "127.0.0.3"
	})

	It("returns successfully without args", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		mock := &PodListerMock{
			PodList: &corev1.PodList{},
		}

		podAddresses, err := WaitForReadyPods(ctx, mock, []string{})
		Expect(err).ToNot(HaveOccurred())
		Expect(podAddresses).To(BeEmpty())
	})

	It("timeout reached when no matching pods are found", func() {
		ctx, cancel := context.WithTimeout(context.Background(), slowDuration)
		defer cancel()

		mock := &PodListerMock{
			PodList: &corev1.PodList{
				Items: []corev1.Pod{irrelevantPod},
			},
		}

		_, err := WaitForReadyPods(ctx, mock, []string{"hello-anyone-out-there"})
		Expect(err).To(HaveOccurred())
	})

	It("timeout reached when only some matching pods are found", func() {
		ctx, cancel := context.WithTimeout(context.Background(), slowDuration)
		defer cancel()

		mock := &PodListerMock{
			PodList: &corev1.PodList{
				Items: []corev1.Pod{
					driverPod,
					clientPod,
					// note: missing 2nd client
				},
			},
		}

		_, err := WaitForReadyPods(ctx, mock, []string{"role=driver", "role=client", "role=client"})
		Expect(err).To(HaveOccurred())
	})

	It("timeout reached when pod does not match all labels", func() {
		ctx, cancel := context.WithTimeout(context.Background(), slowDuration)
		defer cancel()

		mock := &PodListerMock{
			PodList: &corev1.PodList{
				Items: []corev1.Pod{driverPod},
			},
		}

		_, err := WaitForReadyPods(ctx, mock, []string{"role=driver,loadtest=loadtest-1"})
		Expect(err).To(HaveOccurred())
	})

	It("timeout reached when not ready pod matches", func() {
		ctx, cancel := context.WithTimeout(context.Background(), slowDuration)
		defer cancel()

		driverPod.Status.ContainerStatuses[0].Ready = false

		mock := &PodListerMock{
			PodList: &corev1.PodList{
				Items: []corev1.Pod{driverPod},
			},
		}

		_, err := WaitForReadyPods(ctx, mock, []string{"role=driver"})
		Expect(err).To(HaveOccurred())
	})

	It("does not match the same pod twice", func() {
		ctx, cancel := context.WithTimeout(context.Background(), slowDuration)
		defer cancel()

		mock := &PodListerMock{
			PodList: &corev1.PodList{
				Items: []corev1.Pod{
					clientPod,
				},
			},
		}

		_, err := WaitForReadyPods(ctx, mock, []string{
			"role=client",
			"role=client",
		})
		Expect(err).To(HaveOccurred())
	})

	It("returns successfully when all matching pods are found", func() {
		ctx, cancel := context.WithTimeout(context.Background(), slowDuration)
		defer cancel()

		client2Pod := newTestPod("client")
		client2Pod.Name = "client-2"
		client2Pod.Status.PodIP = "127.0.0.4"

		mock := &PodListerMock{
			PodList: &corev1.PodList{
				Items: []corev1.Pod{
					driverPod,
					serverPod,
					clientPod,
					client2Pod,
				},
			},
		}

		podAddresses, err := WaitForReadyPods(ctx, mock, []string{
			"role=driver",
			"role=server",
			"role=client",
			"role=client",
		})
		Expect(err).ToNot(HaveOccurred())
		Expect(podAddresses).To(Equal([]string{
			fmt.Sprintf("%s:%d", driverPod.Status.PodIP, DefaultDriverPort),
			fmt.Sprintf("%s:%d", serverPod.Status.PodIP, DefaultDriverPort),
			fmt.Sprintf("%s:%d", clientPod.Status.PodIP, DefaultDriverPort),
			fmt.Sprintf("%s:%d", client2Pod.Status.PodIP, DefaultDriverPort),
		}))
	})

	It("returns with correct ports for matching pods", func() {
		ctx, cancel := context.WithTimeout(context.Background(), slowDuration)
		defer cancel()

		var customPort int32 = 9542
		client2Pod := newTestPod("client")
		client2Pod.Name = "client-2"
		client2Pod.Spec.Containers[0].Ports[0].ContainerPort = customPort

		mock := &PodListerMock{
			PodList: &corev1.PodList{
				Items: []corev1.Pod{
					clientPod,
					client2Pod,
				},
			},
		}

		podAddresses, err := WaitForReadyPods(ctx, mock, []string{
			"role=client",
			"role=client",
		})
		Expect(err).ToNot(HaveOccurred())
		Expect(podAddresses).To(Equal([]string{
			fmt.Sprintf("%s:%d", clientPod.Status.PodIP, DefaultDriverPort),
			fmt.Sprintf("%s:%d", client2Pod.Status.PodIP, customPort),
		}))
	})

	It("returns error if timeout exceeded", func() {
		ctx, cancel := context.WithTimeout(context.Background(), fastDuration)
		defer cancel()

		mock := &PodListerMock{
			SleepDuration: slowDuration,
			PodList:       &corev1.PodList{},
		}

		_, err := WaitForReadyPods(ctx, mock, []string{"example"})
		Expect(err).To(HaveOccurred())
	})
})

type PodListerMock struct {
	PodList       *corev1.PodList
	SleepDuration time.Duration
	Error         error
	invocation    int
}

func (plm *PodListerMock) List(opts metav1.ListOptions) (*corev1.PodList, error) {
	time.Sleep(plm.SleepDuration)

	if plm.Error != nil {
		return nil, plm.Error
	}

	return plm.PodList, nil
}
<commit_msg>Fix label selector in test case<commit_after>/*
Copyright 2020 gRPC authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"context"
	"fmt"
	"time"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("WaitForReadyPods", func() {
	var fastDuration time.Duration
	var slowDuration time.Duration

	var irrelevantPod corev1.Pod
	var driverPod corev1.Pod
	var serverPod corev1.Pod
	var clientPod corev1.Pod

	BeforeEach(func() {
		fastDuration = 1 * time.Millisecond * timeMultiplier
		slowDuration = 100 * time.Millisecond * timeMultiplier

		irrelevantPod = corev1.Pod{}

		driverPod = newTestPod("driver")
		driverPod.Spec.Containers[0].Ports = nil

		serverPod = newTestPod("server")
		serverPod.Status.PodIP = "127.0.0.2"

		clientPod = newTestPod("client")
		clientPod.Status.PodIP = "127.0.0.3"
	})

	It("returns successfully without args", func() {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		mock := &PodListerMock{
			PodList: &corev1.PodList{},
		}

		podAddresses, err := WaitForReadyPods(ctx, mock, []string{})
		Expect(err).ToNot(HaveOccurred())
		Expect(podAddresses).To(BeEmpty())
	})

	It("timeout reached when no matching pods are found", func() {
		ctx, cancel := context.WithTimeout(context.Background(), slowDuration)
		defer cancel()

		mock := &PodListerMock{
			PodList: &corev1.PodList{
				Items: []corev1.Pod{irrelevantPod},
			},
		}

		_, err := WaitForReadyPods(ctx, mock, []string{"hello=anyone-out-there"})
		Expect(err).To(HaveOccurred())
	})

	It("timeout reached when only some matching pods are found", func() {
		ctx, cancel := context.WithTimeout(context.Background(), slowDuration)
		defer cancel()

		mock := &PodListerMock{
			PodList: &corev1.PodList{
				Items: []corev1.Pod{
					driverPod,
					clientPod,
					// note: missing 2nd client
				},
			},
		}

		_, err := WaitForReadyPods(ctx, mock, []string{"role=driver", "role=client", "role=client"})
		Expect(err).To(HaveOccurred())
	})

	It("timeout reached when pod does not match all labels", func() {
		ctx, cancel := context.WithTimeout(context.Background(), slowDuration)
		defer cancel()

		mock := &PodListerMock{
			PodList: &corev1.PodList{
				Items: []corev1.Pod{driverPod},
			},
		}

		_, err := WaitForReadyPods(ctx, mock, []string{"role=driver,loadtest=loadtest-1"})
		Expect(err).To(HaveOccurred())
	})

	It("timeout reached when not ready pod matches", func() {
		ctx, cancel := context.WithTimeout(context.Background(), slowDuration)
		defer cancel()

		driverPod.Status.ContainerStatuses[0].Ready = false

		mock := &PodListerMock{
			PodList: &corev1.PodList{
				Items: []corev1.Pod{driverPod},
			},
		}

		_, err := WaitForReadyPods(ctx, mock, []string{"role=driver"})
		Expect(err).To(HaveOccurred())
	})

	It("does not match the same pod twice", func() {
		ctx, cancel := context.WithTimeout(context.Background(), slowDuration)
		defer cancel()

		mock := &PodListerMock{
			PodList: &corev1.PodList{
				Items: []corev1.Pod{
					clientPod,
				},
			},
		}

		_, err := WaitForReadyPods(ctx, mock, []string{
			"role=client",
			"role=client",
		})
		Expect(err).To(HaveOccurred())
	})

	It("returns successfully when all matching pods are found", func() {
		ctx, cancel := context.WithTimeout(context.Background(), slowDuration)
		defer cancel()

		client2Pod := newTestPod("client")
		client2Pod.Name = "client-2"
		client2Pod.Status.PodIP = "127.0.0.4"

		mock := &PodListerMock{
			PodList: &corev1.PodList{
				Items: []corev1.Pod{
					driverPod,
					serverPod,
					clientPod,
					client2Pod,
				},
			},
		}

		podAddresses, err := WaitForReadyPods(ctx, mock, []string{
			"role=driver",
			"role=server",
			"role=client",
			"role=client",
		})
		Expect(err).ToNot(HaveOccurred())
		Expect(podAddresses).To(Equal([]string{
			fmt.Sprintf("%s:%d", driverPod.Status.PodIP, DefaultDriverPort),
			fmt.Sprintf("%s:%d", serverPod.Status.PodIP, DefaultDriverPort),
			fmt.Sprintf("%s:%d", clientPod.Status.PodIP, DefaultDriverPort),
			fmt.Sprintf("%s:%d", client2Pod.Status.PodIP, DefaultDriverPort),
		}))
	})

	It("returns with correct ports for matching pods", func() {
		ctx, cancel := context.WithTimeout(context.Background(), slowDuration)
		defer cancel()

		var customPort int32 = 9542
		client2Pod := newTestPod("client")
		client2Pod.Name = "client-2"
		client2Pod.Spec.Containers[0].Ports[0].ContainerPort = customPort

		mock := &PodListerMock{
			PodList: &corev1.PodList{
				Items: []corev1.Pod{
					clientPod,
					client2Pod,
				},
			},
		}

		podAddresses, err := WaitForReadyPods(ctx, mock, []string{
			"role=client",
			"role=client",
		})
		Expect(err).ToNot(HaveOccurred())
		Expect(podAddresses).To(Equal([]string{
			fmt.Sprintf("%s:%d", clientPod.Status.PodIP, DefaultDriverPort),
			fmt.Sprintf("%s:%d", client2Pod.Status.PodIP, customPort),
		}))
	})

	It("returns error if timeout exceeded", func() {
		ctx, cancel := context.WithTimeout(context.Background(), fastDuration)
		defer cancel()

		mock := &PodListerMock{
			SleepDuration: slowDuration,
			PodList:       &corev1.PodList{},
		}

		_, err := WaitForReadyPods(ctx, mock, []string{"example"})
		Expect(err).To(HaveOccurred())
	})
})

type PodListerMock struct {
	PodList       *corev1.PodList
	SleepDuration time.Duration
	Error         error
	invocation    int
}

func (plm *PodListerMock) List(opts metav1.ListOptions) (*corev1.PodList, error) {
	time.Sleep(plm.SleepDuration)

	if plm.Error != nil {
		return nil, plm.Error
	}

	return plm.PodList, nil
}
<|endoftext|>


<commit_before>package cluster

import (
	"errors"
	"fmt"
	"io"
	"sync"

	"github.com/citadel/citadel"
)

var (
	ErrEngineNotConnected = errors.New("engine is not connected to docker's REST API")
)

type Cluster struct {
	mux sync.Mutex

	engines         map[string]*citadel.Engine
	schedulers      map[string]citadel.Scheduler
	resourceManager citadel.ResourceManager
}

func New(manager citadel.ResourceManager, engines ...*citadel.Engine) (*Cluster, error) {
	c := &Cluster{
		engines:         make(map[string]*citadel.Engine),
		schedulers:      make(map[string]citadel.Scheduler),
		resourceManager: manager,
	}

	for _, e := range engines {
		if !e.IsConnected() {
			return nil, ErrEngineNotConnected
		}

		c.engines[e.ID] = e
	}

	return c, nil
}

func (c *Cluster) Events(handler citadel.EventHandler) error {
	for _, e := range c.engines {
		if err := e.Events(handler); err != nil {
			return err
		}
	}

	return nil
}

func (c *Cluster) RegisterScheduler(tpe string, s citadel.Scheduler) error {
	c.mux.Lock()
	defer c.mux.Unlock()

	c.schedulers[tpe] = s

	return nil
}

func (c *Cluster) AddEngine(e *citadel.Engine) error {
	c.mux.Lock()
	defer c.mux.Unlock()

	c.engines[e.ID] = e

	return nil
}

func (c *Cluster) RemoveEngine(e *citadel.Engine) error {
	c.mux.Lock()
	defer c.mux.Unlock()

	delete(c.engines, e.ID)

	return nil
}

// ListContainers returns all the containers running in the cluster
func (c *Cluster) ListContainers(all bool) []*citadel.Container {
	out := []*citadel.Container{}

	for _, e := range c.engines {
		containers, _ := e.ListContainers(all)

		out = append(out, containers...)
	}

	return out
}

func (c *Cluster) Kill(container *citadel.Container, sig int) error {
	c.mux.Lock()
	defer c.mux.Unlock()

	engine := c.engines[container.Engine.ID]
	if engine == nil {
		return fmt.Errorf("engine with id %s is not in cluster", container.Engine.ID)
	}

	return engine.Kill(container, sig)
}

func (c *Cluster) Logs(container *citadel.Container, stdout bool, stderr bool) (io.ReadCloser, error) {
	engine := c.engines[container.Engine.ID]
	if engine == nil {
		return nil, fmt.Errorf("engine with id %s is not in cluster", container.Engine.ID)
	}

	return engine.Logs(container, stdout, stderr)
}

func (c *Cluster) Stop(container *citadel.Container) error {
	c.mux.Lock()
	defer c.mux.Unlock()

	engine := c.engines[container.Engine.ID]
	if engine == nil {
		return fmt.Errorf("engine with id %s is not in cluster", container.Engine.ID)
	}

	return engine.Stop(container)
}

func (c *Cluster) Restart(container *citadel.Container, timeout int) error {
	c.mux.Lock()
	defer c.mux.Unlock()

	engine := c.engines[container.Engine.ID]
	if engine == nil {
		return fmt.Errorf("engine with id %s is not in cluster", container.Engine.ID)
	}

	return engine.Restart(container, timeout)
}

func (c *Cluster) Remove(container *citadel.Container) error {
	c.mux.Lock()
	defer c.mux.Unlock()

	engine := c.engines[container.Engine.ID]
	if engine == nil {
		return fmt.Errorf("engine with id %s is not in cluster", container.Engine.ID)
	}

	return engine.Remove(container)
}

func (c *Cluster) Start(image *citadel.Image, pull bool) (*citadel.Container, error) {
	c.mux.Lock()
	defer c.mux.Unlock()

	var (
		accepted  = []*citadel.EngineSnapshot{}
		scheduler = c.schedulers[image.Type]
	)

	if scheduler == nil {
		return nil, fmt.Errorf("no scheduler for type %s", image.Type)
	}

	for _, e := range c.engines {
		canrun, err := scheduler.Schedule(image, e)
		if err != nil {
			return nil, err
		}

		if canrun {
			containers, err := e.ListContainers(false)
			if err != nil {
				return nil, err
			}

			var cpus, memory float64
			for _, con := range containers {
				cpus += con.Image.Cpus
				memory += con.Image.Memory
			}

			accepted = append(accepted, &citadel.EngineSnapshot{
				ID:             e.ID,
				ReservedCpus:   cpus,
				ReservedMemory: memory,
				Cpus:           e.Cpus,
				Memory:         e.Memory,
			})
		}
	}

	if len(accepted) == 0 {
		return nil, fmt.Errorf("no eligible engines to run image")
	}

	container := &citadel.Container{
		Image: image,
		Name:  image.ContainerName,
	}

	s, err := c.resourceManager.PlaceContainer(container, accepted)
	if err != nil {
		return nil, err
	}

	engine := c.engines[s.ID]

	if err := engine.Start(container, pull); err != nil {
		return nil, err
	}

	return container, nil
}

// Engines returns the engines registered in the cluster
func (c *Cluster) Engines() []*citadel.Engine {
	out := []*citadel.Engine{}

	for _, e := range c.engines {
		out = append(out, e)
	}

	return out
}

// Info returns information about the cluster
func (c *Cluster) ClusterInfo() (*citadel.ClusterInfo, error) {
	containerCount := 0
	imageCount := 0
	engineCount := len(c.engines)
	totalCpu := 0.0
	totalMemory := 0.0
	reservedCpus := 0.0
	reservedMemory := 0.0
	for _, e := range c.engines {
		c, err := e.ListContainers(false)
		if err != nil {
			return nil, err
		}
		for _, cnt := range c {
			reservedCpus += cnt.Image.Cpus
			reservedMemory += cnt.Image.Memory
		}
		i, err := e.ListImages()
		if err != nil {
			return nil, err
		}
		containerCount += len(c)
		imageCount += len(i)
		totalCpu += e.Cpus
		totalMemory += e.Memory
	}

	return &citadel.ClusterInfo{
		Cpus:           totalCpu,
		Memory:         totalMemory,
		ContainerCount: containerCount,
		ImageCount:     imageCount,
		EngineCount:    engineCount,
		ReservedCpus:   reservedCpus,
		ReservedMemory: reservedMemory,
	}, nil
}

// Close signals to the cluster that no other actions will be applied
func (c *Cluster) Close() error {
	return nil
}
<commit_msg>skip engines that are unavailable for info<commit_after>package cluster

import (
	"errors"
	"fmt"
	"io"
	"sync"

	"github.com/citadel/citadel"
)

var (
	ErrEngineNotConnected = errors.New("engine is not connected to docker's REST API")
)

type Cluster struct {
	mux sync.Mutex

	engines         map[string]*citadel.Engine
	schedulers      map[string]citadel.Scheduler
	resourceManager citadel.ResourceManager
}

func New(manager citadel.ResourceManager, engines ...*citadel.Engine) (*Cluster, error) {
	c := &Cluster{
		engines:         make(map[string]*citadel.Engine),
		schedulers:      make(map[string]citadel.Scheduler),
		resourceManager: manager,
	}

	for _, e := range engines {
		if !e.IsConnected() {
			return nil, ErrEngineNotConnected
		}

		c.engines[e.ID] = e
	}

	return c, nil
}

func (c *Cluster) Events(handler citadel.EventHandler) error {
	for _, e := range c.engines {
		if err := e.Events(handler); err != nil {
			return err
		}
	}

	return nil
}

func (c *Cluster) RegisterScheduler(tpe string, s citadel.Scheduler) error {
	c.mux.Lock()
	defer c.mux.Unlock()

	c.schedulers[tpe] = s

	return nil
}

func (c *Cluster) AddEngine(e *citadel.Engine) error {
	c.mux.Lock()
	defer c.mux.Unlock()

	c.engines[e.ID] = e

	return nil
}

func (c *Cluster) RemoveEngine(e *citadel.Engine) error {
	c.mux.Lock()
	defer c.mux.Unlock()

	delete(c.engines, e.ID)

	return nil
}

// ListContainers returns all the containers running in the cluster
func (c *Cluster) ListContainers(all bool) []*citadel.Container {
	out := []*citadel.Container{}

	for _, e := range c.engines {
		containers, _ := e.ListContainers(all)

		out = append(out, containers...)
	}

	return out
}

func (c *Cluster) Kill(container *citadel.Container, sig int) error {
	c.mux.Lock()
	defer c.mux.Unlock()

	engine := c.engines[container.Engine.ID]
	if engine == nil {
		return fmt.Errorf("engine with id %s is not in cluster", container.Engine.ID)
	}

	return engine.Kill(container, sig)
}

func (c *Cluster) Logs(container *citadel.Container, stdout bool, stderr bool) (io.ReadCloser, error) {
	engine := c.engines[container.Engine.ID]
	if engine == nil {
		return nil, fmt.Errorf("engine with id %s is not in cluster", container.Engine.ID)
	}

	return engine.Logs(container, stdout, stderr)
}

func (c *Cluster) Stop(container *citadel.Container) error {
	c.mux.Lock()
	defer c.mux.Unlock()

	engine := c.engines[container.Engine.ID]
	if engine == nil {
		return fmt.Errorf("engine with id %s is not in cluster", container.Engine.ID)
	}

	return engine.Stop(container)
}

func (c *Cluster) Restart(container *citadel.Container, timeout int) error {
	c.mux.Lock()
	defer c.mux.Unlock()

	engine := c.engines[container.Engine.ID]
	if engine == nil {
		return fmt.Errorf("engine with id %s is not in cluster", container.Engine.ID)
	}

	return engine.Restart(container, timeout)
}

func (c *Cluster) Remove(container *citadel.Container) error {
	c.mux.Lock()
	defer c.mux.Unlock()

	engine := c.engines[container.Engine.ID]
	if engine == nil {
		return fmt.Errorf("engine with id %s is not in cluster", container.Engine.ID)
	}

	return engine.Remove(container)
}

func (c *Cluster) Start(image *citadel.Image, pull bool) (*citadel.Container, error) {
	c.mux.Lock()
	defer c.mux.Unlock()

	var (
		accepted  = []*citadel.EngineSnapshot{}
		scheduler = c.schedulers[image.Type]
	)

	if scheduler == nil {
		return nil, fmt.Errorf("no scheduler for type %s", image.Type)
	}

	for _, e := range c.engines {
		canrun, err := scheduler.Schedule(image, e)
		if err != nil {
			return nil, err
		}

		if canrun {
			containers, err := e.ListContainers(false)
			if err != nil {
				return nil, err
			}

			var cpus, memory float64
			for _, con := range containers {
				cpus += con.Image.Cpus
				memory += con.Image.Memory
			}

			accepted = append(accepted, &citadel.EngineSnapshot{
				ID:             e.ID,
				ReservedCpus:   cpus,
				ReservedMemory: memory,
				Cpus:           e.Cpus,
				Memory:         e.Memory,
			})
		}
	}

	if len(accepted) == 0 {
		return nil, fmt.Errorf("no eligible engines to run image")
	}

	container := &citadel.Container{
		Image: image,
		Name:  image.ContainerName,
	}

	s, err := c.resourceManager.PlaceContainer(container, accepted)
	if err != nil {
		return nil, err
	}

	engine := c.engines[s.ID]

	if err := engine.Start(container, pull); err != nil {
		return nil, err
	}

	return container, nil
}

// Engines returns the engines registered in the cluster
func (c *Cluster) Engines() []*citadel.Engine {
	out := []*citadel.Engine{}

	for _, e := range c.engines {
		out = append(out, e)
	}

	return out
}

// Info returns information about the cluster
func (c *Cluster) ClusterInfo() *citadel.ClusterInfo {
	containerCount := 0
	imageCount := 0
	engineCount := len(c.engines)
	totalCpu := 0.0
	totalMemory := 0.0
	reservedCpus := 0.0
	reservedMemory := 0.0
	for _, e := range c.engines {
		c, err := e.ListContainers(false)
		if err != nil {
			// skip engines that are not available
			continue
		}
		for _, cnt := range c {
			reservedCpus += cnt.Image.Cpus
			reservedMemory += cnt.Image.Memory
		}
		i, err := e.ListImages()
		if err != nil {
			// skip engines that are not available
			continue
		}
		containerCount += len(c)
		imageCount += len(i)
		totalCpu += e.Cpus
		totalMemory += e.Memory
	}

	return &citadel.ClusterInfo{
		Cpus:           totalCpu,
		Memory:         totalMemory,
		ContainerCount: containerCount,
		ImageCount:     imageCount,
		EngineCount:    engineCount,
		ReservedCpus:   reservedCpus,
		ReservedMemory: reservedMemory,
	}
}

// Close signals to the cluster that no other actions will be applied
func (c *Cluster) Close() error {
	return nil
}
<|endoftext|>


<commit_before>package mysql

import (
	"database/sql"
	"errors"
	"strings"
	"sync"

	_ "github.com/go-sql-driver/mysql"

	"github.com/henrylee2cn/pholcus/common/util"
	"github.com/henrylee2cn/pholcus/config"
	"github.com/henrylee2cn/pholcus/logs"
)

/************************ Mysql 输出 ***************************/
//sql转换结构体
type MyTable struct {
	tableName        string
	columnNames      [][2]string // 标题字段
	rows             [][]string  // 多行数据
	sqlCode          string
	customPrimaryKey bool
	size             int //内容大小的近似值
}

var (
	db                 *sql.DB
	err                error
	maxConnChan        = make(chan bool, config.MYSQL_CONN_CAP) //最大执行数限制
	max_allowed_packet = config.MYSQL_MAX_ALLOWED_PACKET - 1024
	lock               sync.RWMutex
)

func DB() (*sql.DB, error) {
	return db, err
}

func Refresh() {
	lock.Lock()
	defer lock.Unlock()
	db, err = sql.Open("mysql", config.MYSQL_CONN_STR+"/"+config.DB_NAME+"?charset=utf8")
	if err != nil {
		logs.Log.Error("Mysql：%v\n", err)
		return
	}
	db.SetMaxOpenConns(config.MYSQL_CONN_CAP)
	db.SetMaxIdleConns(config.MYSQL_CONN_CAP)
	if err = db.Ping(); err != nil {
		logs.Log.Error("Mysql：%v\n", err)
	}
}

func New() *MyTable {
	return &MyTable{}
}

//设置表名
func (self *MyTable) SetTableName(name string) *MyTable {
	self.tableName = name
	return self
}

//设置表单列
func (self *MyTable) AddColumn(names ...string) *MyTable {
	for _, name := range names {
		name = strings.Trim(name, " ")
		idx := strings.Index(name, " ")
		self.columnNames = append(self.columnNames, [2]string{string(name[:idx]), string(name[idx+1:])})
	}
	return self
}

//设置主键的语句（可选）
func (self *MyTable) CustomPrimaryKey(primaryKeyCode string) *MyTable {
	self.AddColumn(primaryKeyCode)
	self.customPrimaryKey = true
	return self
}

//生成"创建表单"的语句，执行前须保证SetTableName()、AddColumn()已经执行
func (self *MyTable) Create() error {
	if len(self.columnNames) == 0 {
		return errors.New("Column can not be empty")
	}
	self.sqlCode = `create table if not exists ` + "`" + self.tableName + "` ("
	if !self.customPrimaryKey {
		self.sqlCode += `id int(12) not null primary key auto_increment,`
	}
	for _, title := range self.columnNames {
		self.sqlCode += "`" + title[0] + "`" + ` ` + title[1] + `,`
	}
	self.sqlCode = string(self.sqlCode[:len(self.sqlCode)-1])
	self.sqlCode += `) default charset=utf8;`

	maxConnChan <- true
	defer func() {
		<-maxConnChan
	}()
	lock.RLock()
	stmt, err := db.Prepare(self.sqlCode)
	lock.RUnlock()
	if err != nil {
		return err
	}
	_, err = stmt.Exec()
	return err
}

//清空表单，执行前须保证SetTableName()已经执行
func (self *MyTable) Truncate() error {
	maxConnChan <- true
	defer func() {
		<-maxConnChan
	}()
	lock.RLock()
	stmt, err := db.Prepare(`TRUNCATE TABLE ` + "`" + self.tableName + "`")
	lock.RUnlock()
	if err != nil {
		return err
	}
	_, err = stmt.Exec()
	return err
}

//设置插入的1行数据
func (self *MyTable) addRow(value []string) *MyTable {
	self.rows = append(self.rows, value)
	return self
}

//智能插入数据，每次1行
func (self *MyTable) AutoInsert(value []string) *MyTable {
	var nsize int
	for _, v := range value {
		nsize += len(v)
	}
	if nsize > max_allowed_packet {
		logs.Log.Error("%v", "packet for query is too large. Try adjusting the 'maxallowedpacket'variable in the 'config.ini'")
		return self
	}
	self.size += nsize
	if self.size > max_allowed_packet {
		util.CheckErr(self.FlushInsert())
		return self.AutoInsert(value)
	}
	return self.addRow(value)
}

//向sqlCode添加"插入数据"的语句，执行前须保证Create()、AutoInsert()已经执行
//insert into table1(field1,field2) values(rows[0]),(rows[1])...
func (self *MyTable) FlushInsert() error {
	if len(self.rows) == 0 {
		return nil
	}

	self.sqlCode = `insert into ` + "`" + self.tableName + "`" + `(`
	if len(self.columnNames) != 0 {
		for _, v := range self.columnNames {
			self.sqlCode += "`" + v[0] + "`,"
		}
		self.sqlCode = self.sqlCode[:len(self.sqlCode)-1] + `)values`
	}
	for _, row := range self.rows {
		self.sqlCode += `(`
		for _, v := range row {
			v = strings.Replace(v, `\`, `\\`, -1)
			v = strings.Replace(v, `"`, `\"`, -1)
			v = strings.Replace(v, `'`, `\'`, -1)
			self.sqlCode += `"` + v + `",`
		}
		self.sqlCode = self.sqlCode[:len(self.sqlCode)-1] + `),`
	}
	self.sqlCode = self.sqlCode[:len(self.sqlCode)-1] + `;`

	defer func() {
		// 清空临时数据
		self.rows = [][]string{}
		self.size = 0
		self.sqlCode = ""
	}()

	maxConnChan <- true
	defer func() {
		<-maxConnChan
	}()
	lock.RLock()
	stmt, err := db.Prepare(self.sqlCode)
	lock.RUnlock()
	if err != nil {
		return err
	}
	_, err = stmt.Exec()
	return err
}

// 获取全部数据
func (self *MyTable) SelectAll() (*sql.Rows, error) {
	if self.tableName == "" {
		return nil, errors.New("表名不能为空")
	}
	self.sqlCode = `select * from ` + "`" + self.tableName + "`" + `;`

	maxConnChan <- true
	defer func() {
		<-maxConnChan
	}()
	lock.RLock()
	defer lock.RUnlock()
	return db.Query(self.sqlCode)
}
<commit_msg>修复 mysql 的"Can't create more than max_prepared_stmt_count"bug<commit_after>package mysql

import (
	"database/sql"
	"errors"
	"strings"
	"sync"

	_ "github.com/go-sql-driver/mysql"

	"github.com/henrylee2cn/pholcus/common/util"
	"github.com/henrylee2cn/pholcus/config"
	"github.com/henrylee2cn/pholcus/logs"
)

/************************ Mysql 输出 ***************************/
//sql转换结构体
type MyTable struct {
	tableName        string
	columnNames      [][2]string // 标题字段
	rows             [][]string  // 多行数据
	sqlCode          string
	customPrimaryKey bool
	size             int //内容大小的近似值
}

var (
	db                 *sql.DB
	err                error
	maxConnChan        = make(chan bool, config.MYSQL_CONN_CAP) //最大执行数限制
	max_allowed_packet = config.MYSQL_MAX_ALLOWED_PACKET - 1024
	lock               sync.RWMutex
)

func DB() (*sql.DB, error) {
	return db, err
}

func Refresh() {
	lock.Lock()
	defer lock.Unlock()
	db, err = sql.Open("mysql", config.MYSQL_CONN_STR+"/"+config.DB_NAME+"?charset=utf8")
	if err != nil {
		logs.Log.Error("Mysql：%v\n", err)
		return
	}
	db.SetMaxOpenConns(config.MYSQL_CONN_CAP)
	db.SetMaxIdleConns(config.MYSQL_CONN_CAP)
	if err = db.Ping(); err != nil {
		logs.Log.Error("Mysql：%v\n", err)
	}
}

func New() *MyTable {
	return &MyTable{}
}

//设置表名
func (self *MyTable) SetTableName(name string) *MyTable {
	self.tableName = name
	return self
}

//设置表单列
func (self *MyTable) AddColumn(names ...string) *MyTable {
	for _, name := range names {
		name = strings.Trim(name, " ")
		idx := strings.Index(name, " ")
		self.columnNames = append(self.columnNames, [2]string{string(name[:idx]), string(name[idx+1:])})
	}
	return self
}

//设置主键的语句（可选）
func (self *MyTable) CustomPrimaryKey(primaryKeyCode string) *MyTable {
	self.AddColumn(primaryKeyCode)
	self.customPrimaryKey = true
	return self
}

//生成"创建表单"的语句，执行前须保证SetTableName()、AddColumn()已经执行
func (self *MyTable) Create() error {
	if len(self.columnNames) == 0 {
		return errors.New("Column can not be empty")
	}
	self.sqlCode = `create table if not exists ` + "`" + self.tableName + "` ("
	if !self.customPrimaryKey {
		self.sqlCode += `id int(12) not null primary key auto_increment,`
	}
	for _, title := range self.columnNames {
		self.sqlCode += "`" + title[0] + "`" + ` ` + title[1] + `,`
	}
	self.sqlCode = string(self.sqlCode[:len(self.sqlCode)-1])
	self.sqlCode += `) default charset=utf8;`

	maxConnChan <- true
	defer func() {
		<-maxConnChan
	}()
	lock.RLock()
	stmt, err := db.Prepare(self.sqlCode)
	lock.RUnlock()
	if err != nil {
		return err
	}
	_, err = stmt.Exec()
	stmt.Close()
	return err
}

//清空表单，执行前须保证SetTableName()已经执行
func (self *MyTable) Truncate() error {
	maxConnChan <- true
	defer func() {
		<-maxConnChan
	}()
	lock.RLock()
	stmt, err := db.Prepare(`TRUNCATE TABLE ` + "`" + self.tableName + "`")
	lock.RUnlock()
	if err != nil {
		return err
	}
	_, err = stmt.Exec()
	stmt.Close()
	return err
}

//设置插入的1行数据
func (self *MyTable) addRow(value []string) *MyTable {
	self.rows = append(self.rows, value)
	return self
}

//智能插入数据，每次1行
func (self *MyTable) AutoInsert(value []string) *MyTable {
	var nsize int
	for _, v := range value {
		nsize += len(v)
	}
	if nsize > max_allowed_packet {
		logs.Log.Error("%v", "packet for query is too large. Try adjusting the 'maxallowedpacket'variable in the 'config.ini'")
		return self
	}
	self.size += nsize
	if self.size > max_allowed_packet {
		util.CheckErr(self.FlushInsert())
		return self.AutoInsert(value)
	}
	return self.addRow(value)
}

//向sqlCode添加"插入数据"的语句，执行前须保证Create()、AutoInsert()已经执行
//insert into table1(field1,field2) values(rows[0]),(rows[1])...
func (self *MyTable) FlushInsert() error {
	if len(self.rows) == 0 {
		return nil
	}

	self.sqlCode = `insert into ` + "`" + self.tableName + "`" + `(`
	if len(self.columnNames) != 0 {
		for _, v := range self.columnNames {
			self.sqlCode += "`" + v[0] + "`,"
		}
		self.sqlCode = self.sqlCode[:len(self.sqlCode)-1] + `)values`
	}
	for _, row := range self.rows {
		self.sqlCode += `(`
		for _, v := range row {
			v = strings.Replace(v, `\`, `\\`, -1)
			v = strings.Replace(v, `"`, `\"`, -1)
			v = strings.Replace(v, `'`, `\'`, -1)
			self.sqlCode += `"` + v + `",`
		}
		self.sqlCode = self.sqlCode[:len(self.sqlCode)-1] + `),`
	}
	self.sqlCode = self.sqlCode[:len(self.sqlCode)-1] + `;`

	defer func() {
		// 清空临时数据
		self.rows = [][]string{}
		self.size = 0
		self.sqlCode = ""
	}()

	maxConnChan <- true
	defer func() {
		<-maxConnChan
	}()
	lock.RLock()
	stmt, err := db.Prepare(self.sqlCode)
	lock.RUnlock()
	if err != nil {
		return err
	}
	_, err = stmt.Exec()
	stmt.Close()
	return err
}

// 获取全部数据
func (self *MyTable) SelectAll() (*sql.Rows, error) {
	if self.tableName == "" {
		return nil, errors.New("表名不能为空")
	}
	self.sqlCode = `select * from ` + "`" + self.tableName + "`" + `;`

	maxConnChan <- true
	defer func() {
		<-maxConnChan
	}()
	lock.RLock()
	defer lock.RUnlock()
	return db.Query(self.sqlCode)
}
<|endoftext|>


<commit_before>package contractcourt

import (
	"bytes"
	"crypto/sha256"
	"testing"
	"time"

	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/wire"
	"github.com/lightningnetwork/lnd/chainntnfs"
	"github.com/lightningnetwork/lnd/lnwallet"
	"github.com/lightningnetwork/lnd/lnwire"
)

type mockNotifier struct {
	spendChan chan *chainntnfs.SpendDetail
	epochChan chan *chainntnfs.BlockEpoch
	confChan  chan *chainntnfs.TxConfirmation
}

func (m *mockNotifier) RegisterConfirmationsNtfn(txid *chainhash.Hash, _ []byte, numConfs,
	heightHint uint32) (*chainntnfs.ConfirmationEvent, error) {
	return &chainntnfs.ConfirmationEvent{
		Confirmed: m.confChan,
	}, nil
}

func (m *mockNotifier) RegisterBlockEpochNtfn(
	bestBlock *chainntnfs.BlockEpoch) (*chainntnfs.BlockEpochEvent, error) {

	return &chainntnfs.BlockEpochEvent{
		Epochs: m.epochChan,
		Cancel: func() {},
	}, nil
}

func (m *mockNotifier) Start() error {
	return nil
}

func (m *mockNotifier) Stop() error {
	return nil
}
func (m *mockNotifier) RegisterSpendNtfn(outpoint *wire.OutPoint, _ []byte,
	heightHint uint32) (*chainntnfs.SpendEvent, error) {

	return &chainntnfs.SpendEvent{
		Spend:  m.spendChan,
		Cancel: func() {},
	}, nil
}

// TestChainWatcherRemoteUnilateralClose tests that the chain watcher is able
// to properly detect a normal unilateral close by the remote node using their
// lowest commitment.
func TestChainWatcherRemoteUnilateralClose(t *testing.T) {
	t.Parallel()

	// First, we'll create two channels which already have established a
	// commitment contract between themselves.
	aliceChannel, bobChannel, cleanUp, err := lnwallet.CreateTestChannels()
	if err != nil {
		t.Fatalf("unable to create test channels: %v", err)
	}
	defer cleanUp()

	// With the channels created, we'll now create a chain watcher instance
	// which will be watching for any closes of Alice's channel.
	aliceNotifier := &mockNotifier{
		spendChan: make(chan *chainntnfs.SpendDetail),
	}
	aliceChainWatcher, err := newChainWatcher(chainWatcherConfig{
		chanState: aliceChannel.State(),
		notifier:  aliceNotifier,
		signer:    aliceChannel.Signer,
	})
	if err != nil {
		t.Fatalf("unable to create chain watcher: %v", err)
	}
	err = aliceChainWatcher.Start()
	if err != nil {
		t.Fatalf("unable to start chain watcher: %v", err)
	}
	defer aliceChainWatcher.Stop()

	// We'll request a new channel event subscription from Alice's chain
	// watcher.
	chanEvents := aliceChainWatcher.SubscribeChannelEvents()

	// If we simulate an immediate broadcast of the current commitment by
	// Bob, then the chain watcher should detect this case.
	bobCommit := bobChannel.State().LocalCommitment.CommitTx
	bobTxHash := bobCommit.TxHash()
	bobSpend := &chainntnfs.SpendDetail{
		SpenderTxHash: &bobTxHash,
		SpendingTx:    bobCommit,
	}
	aliceNotifier.spendChan <- bobSpend

	// We should get a new spend event over the remote unilateral close
	// event channel.
	var uniClose *lnwallet.UnilateralCloseSummary
	select {
	case uniClose = <-chanEvents.RemoteUnilateralClosure:
	case <-time.After(time.Second * 15):
		t.Fatalf("didn't receive unilateral close event")
	}

	// The unilateral close should have properly located Alice's output in
	// the commitment transaction.
	if uniClose.CommitResolution == nil {
		t.Fatalf("unable to find alice's commit resolution")
	}
}

// TestChainWatcherRemoteUnilateralClosePendingCommit tests that the chain
// watcher is able to properly detect a unilateral close wherein the remote
// node broadcasts their newly received commitment, without first revoking the
// old one.
func TestChainWatcherRemoteUnilateralClosePendingCommit(t *testing.T) {
	t.Parallel()

	// First, we'll create two channels which already have established a
	// commitment contract between themselves.
	aliceChannel, bobChannel, cleanUp, err := lnwallet.CreateTestChannels()
	if err != nil {
		t.Fatalf("unable to create test channels: %v", err)
	}
	defer cleanUp()

	// With the channels created, we'll now create a chain watcher instance
	// which will be watching for any closes of Alice's channel.
	aliceNotifier := &mockNotifier{
		spendChan: make(chan *chainntnfs.SpendDetail),
	}
	aliceChainWatcher, err := newChainWatcher(chainWatcherConfig{
		chanState: aliceChannel.State(),
		notifier:  aliceNotifier,
		signer:    aliceChannel.Signer,
	})
	if err != nil {
		t.Fatalf("unable to create chain watcher: %v", err)
	}
	if err := aliceChainWatcher.Start(); err != nil {
		t.Fatalf("unable to start chain watcher: %v", err)
	}
	defer aliceChainWatcher.Stop()

	// We'll request a new channel event subscription from Alice's chain
	// watcher.
	chanEvents := aliceChainWatcher.SubscribeChannelEvents()

	// Next, we'll create a fake HTLC just so we can advance Alice's
	// channel state to a new pending commitment on her remote commit chain
	// for Bob.
	htlcAmount := lnwire.NewMSatFromSatoshis(20000)
	preimage := bytes.Repeat([]byte{byte(1)}, 32)
	paymentHash := sha256.Sum256(preimage)
	var returnPreimage [32]byte
	copy(returnPreimage[:], preimage)
	htlc := &lnwire.UpdateAddHTLC{
		ID:          uint64(0),
		PaymentHash: paymentHash,
		Amount:      htlcAmount,
		Expiry:      uint32(5),
	}

	if _, err := aliceChannel.AddHTLC(htlc, nil); err != nil {
		t.Fatalf("alice unable to add htlc: %v", err)
	}
	if _, err := bobChannel.ReceiveHTLC(htlc); err != nil {
		t.Fatalf("bob unable to recv add htlc: %v", err)
	}

	// With the HTLC added, we'll now manually initiate a state transition
	// from Alice to Bob.
	_, _, err = aliceChannel.SignNextCommitment()
	if err != nil {
		t.Fatal(err)
	}

	// At this point, we'll now Bob broadcasting this new pending unrevoked
	// commitment.
	bobPendingCommit, err := aliceChannel.State().RemoteCommitChainTip()
	if err != nil {
		t.Fatal(err)
	}

	// We'll craft a fake spend notification with Bob's actual commitment.
	// The chain watcher should be able to detect that this is a pending
	// commit broadcast based on the state hints in the commitment.
	bobCommit := bobPendingCommit.Commitment.CommitTx
	bobTxHash := bobCommit.TxHash()
	bobSpend := &chainntnfs.SpendDetail{
		SpenderTxHash: &bobTxHash,
		SpendingTx:    bobCommit,
	}
	aliceNotifier.spendChan <- bobSpend

	// We should get a new spend event over the remote unilateral close
	// event channel.
	var uniClose *lnwallet.UnilateralCloseSummary
	select {
	case uniClose = <-chanEvents.RemoteUnilateralClosure:
	case <-time.After(time.Second * 15):
		t.Fatalf("didn't receive unilateral close event")
	}

	// The unilateral close should have properly located Alice's output in
	// the commitment transaction.
	if uniClose.CommitResolution == nil {
		t.Fatalf("unable to find alice's commit resolution")
	}
}
<commit_msg>contractcourt: add new TestChainWatcherDataLossProtect test case<commit_after>package contractcourt

import (
	"bytes"
	"crypto/sha256"
	"math"
	"math/rand"
	"reflect"
	"testing"
	"testing/quick"
	"time"

	"github.com/btcsuite/btcd/chaincfg/chainhash"
	"github.com/btcsuite/btcd/wire"
	"github.com/lightningnetwork/lnd/chainntnfs"
	"github.com/lightningnetwork/lnd/channeldb"
	"github.com/lightningnetwork/lnd/input"
	"github.com/lightningnetwork/lnd/lnwallet"
	"github.com/lightningnetwork/lnd/lnwire"
)

type mockNotifier struct {
	spendChan chan *chainntnfs.SpendDetail
	epochChan chan *chainntnfs.BlockEpoch
	confChan  chan *chainntnfs.TxConfirmation
}

func (m *mockNotifier) RegisterConfirmationsNtfn(txid *chainhash.Hash, _ []byte, numConfs,
	heightHint uint32) (*chainntnfs.ConfirmationEvent, error) {
	return &chainntnfs.ConfirmationEvent{
		Confirmed: m.confChan,
	}, nil
}

func (m *mockNotifier) RegisterBlockEpochNtfn(
	bestBlock *chainntnfs.BlockEpoch) (*chainntnfs.BlockEpochEvent, error) {

	return &chainntnfs.BlockEpochEvent{
		Epochs: m.epochChan,
		Cancel: func() {},
	}, nil
}

func (m *mockNotifier) Start() error {
	return nil
}

func (m *mockNotifier) Stop() error {
	return nil
}
func (m *mockNotifier) RegisterSpendNtfn(outpoint *wire.OutPoint, _ []byte,
	heightHint uint32) (*chainntnfs.SpendEvent, error) {

	return &chainntnfs.SpendEvent{
		Spend:  m.spendChan,
		Cancel: func() {},
	}, nil
}

// TestChainWatcherRemoteUnilateralClose tests that the chain watcher is able
// to properly detect a normal unilateral close by the remote node using their
// lowest commitment.
func TestChainWatcherRemoteUnilateralClose(t *testing.T) {
	t.Parallel()

	// First, we'll create two channels which already have established a
	// commitment contract between themselves.
	aliceChannel, bobChannel, cleanUp, err := lnwallet.CreateTestChannels()
	if err != nil {
		t.Fatalf("unable to create test channels: %v", err)
	}
	defer cleanUp()

	// With the channels created, we'll now create a chain watcher instance
	// which will be watching for any closes of Alice's channel.
	aliceNotifier := &mockNotifier{
		spendChan: make(chan *chainntnfs.SpendDetail),
	}
	aliceChainWatcher, err := newChainWatcher(chainWatcherConfig{
		chanState:           aliceChannel.State(),
		notifier:            aliceNotifier,
		signer:              aliceChannel.Signer,
		extractStateNumHint: lnwallet.GetStateNumHint,
	})
	if err != nil {
		t.Fatalf("unable to create chain watcher: %v", err)
	}
	err = aliceChainWatcher.Start()
	if err != nil {
		t.Fatalf("unable to start chain watcher: %v", err)
	}
	defer aliceChainWatcher.Stop()

	// We'll request a new channel event subscription from Alice's chain
	// watcher.
	chanEvents := aliceChainWatcher.SubscribeChannelEvents()

	// If we simulate an immediate broadcast of the current commitment by
	// Bob, then the chain watcher should detect this case.
	bobCommit := bobChannel.State().LocalCommitment.CommitTx
	bobTxHash := bobCommit.TxHash()
	bobSpend := &chainntnfs.SpendDetail{
		SpenderTxHash: &bobTxHash,
		SpendingTx:    bobCommit,
	}
	aliceNotifier.spendChan <- bobSpend

	// We should get a new spend event over the remote unilateral close
	// event channel.
	var uniClose *lnwallet.UnilateralCloseSummary
	select {
	case uniClose = <-chanEvents.RemoteUnilateralClosure:
	case <-time.After(time.Second * 15):
		t.Fatalf("didn't receive unilateral close event")
	}

	// The unilateral close should have properly located Alice's output in
	// the commitment transaction.
	if uniClose.CommitResolution == nil {
		t.Fatalf("unable to find alice's commit resolution")
	}
}

func addFakeHTLC(t *testing.T, htlcAmount lnwire.MilliSatoshi, id uint64,
	aliceChannel, bobChannel *lnwallet.LightningChannel) {

	preimage := bytes.Repeat([]byte{byte(id)}, 32)
	paymentHash := sha256.Sum256(preimage)
	var returnPreimage [32]byte
	copy(returnPreimage[:], preimage)
	htlc := &lnwire.UpdateAddHTLC{
		ID:          uint64(id),
		PaymentHash: paymentHash,
		Amount:      htlcAmount,
		Expiry:      uint32(5),
	}

	if _, err := aliceChannel.AddHTLC(htlc, nil); err != nil {
		t.Fatalf("alice unable to add htlc: %v", err)
	}
	if _, err := bobChannel.ReceiveHTLC(htlc); err != nil {
		t.Fatalf("bob unable to recv add htlc: %v", err)
	}
}

// TestChainWatcherRemoteUnilateralClosePendingCommit tests that the chain
// watcher is able to properly detect a unilateral close wherein the remote
// node broadcasts their newly received commitment, without first revoking the
// old one.
func TestChainWatcherRemoteUnilateralClosePendingCommit(t *testing.T) {
	t.Parallel()

	// First, we'll create two channels which already have established a
	// commitment contract between themselves.
	aliceChannel, bobChannel, cleanUp, err := lnwallet.CreateTestChannels()
	if err != nil {
		t.Fatalf("unable to create test channels: %v", err)
	}
	defer cleanUp()

	// With the channels created, we'll now create a chain watcher instance
	// which will be watching for any closes of Alice's channel.
	aliceNotifier := &mockNotifier{
		spendChan: make(chan *chainntnfs.SpendDetail),
	}
	aliceChainWatcher, err := newChainWatcher(chainWatcherConfig{
		chanState:           aliceChannel.State(),
		notifier:            aliceNotifier,
		signer:              aliceChannel.Signer,
		extractStateNumHint: lnwallet.GetStateNumHint,
	})
	if err != nil {
		t.Fatalf("unable to create chain watcher: %v", err)
	}
	if err := aliceChainWatcher.Start(); err != nil {
		t.Fatalf("unable to start chain watcher: %v", err)
	}
	defer aliceChainWatcher.Stop()

	// We'll request a new channel event subscription from Alice's chain
	// watcher.
	chanEvents := aliceChainWatcher.SubscribeChannelEvents()

	// Next, we'll create a fake HTLC just so we can advance Alice's
	// channel state to a new pending commitment on her remote commit chain
	// for Bob.
	htlcAmount := lnwire.NewMSatFromSatoshis(20000)
	addFakeHTLC(t, htlcAmount, 0, aliceChannel, bobChannel)

	// With the HTLC added, we'll now manually initiate a state transition
	// from Alice to Bob.
	_, _, err = aliceChannel.SignNextCommitment()
	if err != nil {
		t.Fatal(err)
	}

	// At this point, we'll now Bob broadcasting this new pending unrevoked
	// commitment.
	bobPendingCommit, err := aliceChannel.State().RemoteCommitChainTip()
	if err != nil {
		t.Fatal(err)
	}

	// We'll craft a fake spend notification with Bob's actual commitment.
	// The chain watcher should be able to detect that this is a pending
	// commit broadcast based on the state hints in the commitment.
	bobCommit := bobPendingCommit.Commitment.CommitTx
	bobTxHash := bobCommit.TxHash()
	bobSpend := &chainntnfs.SpendDetail{
		SpenderTxHash: &bobTxHash,
		SpendingTx:    bobCommit,
	}
	aliceNotifier.spendChan <- bobSpend

	// We should get a new spend event over the remote unilateral close
	// event channel.
	var uniClose *lnwallet.UnilateralCloseSummary
	select {
	case uniClose = <-chanEvents.RemoteUnilateralClosure:
	case <-time.After(time.Second * 15):
		t.Fatalf("didn't receive unilateral close event")
	}

	// The unilateral close should have properly located Alice's output in
	// the commitment transaction.
	if uniClose.CommitResolution == nil {
		t.Fatalf("unable to find alice's commit resolution")
	}
}

// dlpTestCase is a speical struct that we'll use to generate randomized test
// cases for the main TestChainWatcherDataLossProtect test. This struct has a
// special Generate method that will generate a random state number, and a
// broadcast state number which is greater than that state number.
type dlpTestCase struct {
	BroadcastStateNum uint8
	NumUpdates        uint8
}

// TestChainWatcherDataLossProtect tests that if we've lost data (and are
// behind the remote node), then we'll properly detect this case and dispatch a
// remote force close using the obtained data loss commitment point.
func TestChainWatcherDataLossProtect(t *testing.T) {
	t.Parallel()

	// dlpScenario is our primary quick check testing function for this
	// test as whole. It ensures that if the remote party broadcasts a
	// commitment that is beyond our best known commitment for them, and
	// they don't have a pending commitment (one we sent but which hasn't
	// been revoked), then we'll properly detect this case, and execute the
	// DLP protocol on our end.
	//
	// broadcastStateNum is the number that we'll trick Alice into thinking
	// was broadcast, while numUpdates is the actual number of updates
	// we'll execute. Both of these will be random 8-bit values generated
	// by testing/quick.
	dlpScenario := func(testCase dlpTestCase) bool {
		// First, we'll create two channels which already have
		// established a commitment contract between themselves.
		aliceChannel, bobChannel, cleanUp, err := lnwallet.CreateTestChannels()
		if err != nil {
			t.Fatalf("unable to create test channels: %v", err)
		}
		defer cleanUp()

		// With the channels created, we'll now create a chain watcher
		// instance which will be watching for any closes of Alice's
		// channel.
		aliceNotifier := &mockNotifier{
			spendChan: make(chan *chainntnfs.SpendDetail),
		}
		aliceChainWatcher, err := newChainWatcher(chainWatcherConfig{
			chanState: aliceChannel.State(),
			notifier:  aliceNotifier,
			signer:    aliceChannel.Signer,
			extractStateNumHint: func(*wire.MsgTx,
				[lnwallet.StateHintSize]byte) uint64 {

				// We'll return the "fake" broadcast commitment
				// number so we can simulate broadcast of an
				// arbitrary state.
				return uint64(testCase.BroadcastStateNum)
			},
		})
		if err != nil {
			t.Fatalf("unable to create chain watcher: %v", err)
		}
		if err := aliceChainWatcher.Start(); err != nil {
			t.Fatalf("unable to start chain watcher: %v", err)
		}
		defer aliceChainWatcher.Stop()

		// Based on the number of random updates for this state, make a
		// new HTLC to add to the commitment, and then lock in a state
		// transition.
		const htlcAmt = 1000
		for i := 0; i < int(testCase.NumUpdates); i++ {
			addFakeHTLC(
				t, 1000, uint64(i), aliceChannel, bobChannel,
			)

			err := lnwallet.ForceStateTransition(
				aliceChannel, bobChannel,
			)
			if err != nil {
				t.Errorf("unable to trigger state "+
					"transition: %v", err)
				return false
			}
		}

		// We'll request a new channel event subscription from Alice's
		// chain watcher so we can be notified of our fake close below.
		chanEvents := aliceChainWatcher.SubscribeChannelEvents()

		// Otherwise, we'll feed in this new state number as a response
		// to the query, and insert the expected DLP commit point.
		dlpPoint := aliceChannel.State().RemoteCurrentRevocation
		err = aliceChannel.State().MarkDataLoss(dlpPoint)
		if err != nil {
			t.Errorf("unable to insert dlp point: %v", err)
			return false
		}

		// Now we'll trigger the channel close event to trigger the
		// scenario.
		bobCommit := bobChannel.State().LocalCommitment.CommitTx
		bobTxHash := bobCommit.TxHash()
		bobSpend := &chainntnfs.SpendDetail{
			SpenderTxHash: &bobTxHash,
			SpendingTx:    bobCommit,
		}
		aliceNotifier.spendChan <- bobSpend

		// We should get a new uni close resolution that indicates we
		// processed the DLP scenario.
		var uniClose *lnwallet.UnilateralCloseSummary
		select {
		case uniClose = <-chanEvents.RemoteUnilateralClosure:
			// If we processed this as a DLP case, then the remote
			// party's commitment should be blank, as we don't have
			// this up to date state.
			blankCommit := channeldb.ChannelCommitment{}
			if uniClose.RemoteCommit.FeePerKw != blankCommit.FeePerKw {
				t.Errorf("DLP path not executed")
				return false
			}

			// The resolution should have also read the DLP point
			// we stored above, and used that to derive their sweep
			// key for this output.
			sweepTweak := input.SingleTweakBytes(
				dlpPoint,
				aliceChannel.State().LocalChanCfg.PaymentBasePoint.PubKey,
			)
			commitResolution := uniClose.CommitResolution
			resolutionTweak := commitResolution.SelfOutputSignDesc.SingleTweak
			if !bytes.Equal(sweepTweak, resolutionTweak) {
				t.Errorf("sweep key mismatch: expected %x got %x",
					sweepTweak, resolutionTweak)
				return false
			}

			return true

		case <-time.After(time.Second * 5):
			t.Errorf("didn't receive unilateral close event")
			return false
		}
	}

	// For our first scenario, we'll ensure that if we're on state 1, and
	// the remote party broadcasts state 2 and we don't have a pending
	// commit for them, then we'll properly detect this as a DLP scenario.
	if !dlpScenario(dlpTestCase{
		BroadcastStateNum: 2,
		NumUpdates:        1,
	}) {
		t.Fatalf("DLP test case failed at state 1!")
	}

	// For the remainder of the tests, we'll perform 10 iterations with
	// random values. We limit this number as set up of each test can take
	// time, and also it doing up to 255 state transitions may cause the
	// test to hang for a long time.
	//
	// TODO(roasbeef): speed up execution
	err := quick.Check(dlpScenario, &quick.Config{
		MaxCount: 10,
		Values: func(v []reflect.Value, rand *rand.Rand) {
			// stateNum will be the random number of state updates
			// we execute during the scenario.
			stateNum := uint8(rand.Int31())

			// From this state number, we'll draw a random number
			// between the state and 255, ensuring that it' at
			// least one state beyond the target stateNum.
			broadcastRange := rand.Int31n(int32(math.MaxUint8 - stateNum))
			broadcastNum := uint8(stateNum + 1 + uint8(broadcastRange))

			testCase := dlpTestCase{
				BroadcastStateNum: broadcastNum,
				NumUpdates:        stateNum,
			}
			v[0] = reflect.ValueOf(testCase)
		},
	})
	if err != nil {
		t.Fatalf("DLP test case failed: %v", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"text/template"

	"github.com/containers/image/types"
	"github.com/cri-o/cri-o/server"
	"github.com/urfave/cli"
)

// NOTE: please propagate any changes to the template to docs/crio.conf.5.md

var commentedConfigTemplate = template.Must(template.New("config").Parse(`
# The CRI-O configuration file specifies all of the available configuration
# options and command-line flags for the crio(8) OCI Kubernetes Container Runtime
# daemon, but in a TOML format that can be more easily modified and versioned.
#
# Please refer to crio.conf(5) for details of all configuration options.

# CRI-O supports partial configuration reload during runtime, which can be
# done by sending SIGHUP to the running process. Currently supported options
# are explicitly mentioned with: 'This option supports live configuration
# reload'.

# CRI-O reads its storage defaults from the containers-storage.conf(5) file
# located at /etc/containers/storage.conf. Modify this storage configuration if
# you want to change the system's defaults. If you want to modify storage just
# for CRI-O, you can change the storage configuration options here.
[crio]

# Path to the "root directory". CRI-O stores all of its data, including
# containers images, in this directory.
#root = "{{ .Root }}"

# Path to the "run directory". CRI-O stores all of its state in this directory.
#runroot = "{{ .RunRoot }}"

# Storage driver used to manage the storage of images and containers. Please
# refer to containers-storage.conf(5) to see all available storage drivers.
#storage_driver = "{{ .Storage }}"

# List to pass options to the storage driver. Please refer to
# containers-storage.conf(5) to see all available storage options.
#storage_option = [
{{ range $opt := .StorageOptions }}{{ printf "#\t%q,\n" $opt }}{{ end }}#]

# If set to false, in-memory locking will be used instead of file-based locking.
# **Deprecated** this option will be removed in the future.
file_locking = {{ .FileLocking }}

# Path to the lock file.
# **Deprecated** this option will be removed in the future.
file_locking_path = "{{ .FileLockingPath }}"

# The crio.api table contains settings for the kubelet/gRPC interface.
[crio.api]

# Path to AF_LOCAL socket on which CRI-O will listen.
listen = "{{ .Listen }}"

# IP address on which the stream server will listen.
stream_address = "{{ .StreamAddress }}"

# The port on which the stream server will listen.
stream_port = "{{ .StreamPort }}"

# Enable encrypted TLS transport of the stream server.
stream_enable_tls = {{ .StreamEnableTLS }}

# Path to the x509 certificate file used to serve the encrypted stream. This
# file can change, and CRI-O will automatically pick up the changes within 5
# minutes.
stream_tls_cert = "{{ .StreamTLSCert }}"

# Path to the key file used to serve the encrypted stream. This file can
# change, and CRI-O will automatically pick up the changes within 5 minutes.
stream_tls_key = "{{ .StreamTLSKey }}"

# Path to the x509 CA(s) file used to verify and authenticate client
# communication with the encrypted stream. This file can change, and CRI-O will
# automatically pick up the changes within 5 minutes.
stream_tls_ca = "{{ .StreamTLSCA }}"

# Maximum grpc send message size in bytes. If not set or <=0, then CRI-O will default to 16 * 1024 * 1024.
grpc_max_send_msg_size = {{ .GRPCMaxSendMsgSize }}

# Maximum grpc receive message size. If not set or <= 0, then CRI-O will default to 16 * 1024 * 1024.
grpc_max_recv_msg_size = {{ .GRPCMaxRecvMsgSize }}

# The crio.runtime table contains settings pertaining to the OCI runtime used
# and options for how to set up and manage the OCI runtime.
[crio.runtime]

# A list of ulimits to be set in containers by default, specified as
# "<ulimit name>=<soft limit>:<hard limit>", for example:
# "nofile=1024:2048"
# If nothing is set here, settings will be inherited from the CRI-O daemon
#default_ulimits = [
{{ range $ulimit := .DefaultUlimits }}{{ printf "#\t%q,\n" $ulimit }}{{ end }}#]

# default_runtime is the _name_ of the OCI runtime to be used as the default.
# The name is matched against the runtimes map below.
default_runtime = "{{ .DefaultRuntime }}"

# If true, the runtime will not use pivot_root, but instead use MS_MOVE.
no_pivot = {{ .NoPivot }}

# Path to the conmon binary, used for monitoring the OCI runtime.
# Will be searched for using $PATH if empty.
conmon = "{{ .Conmon }}"

# Cgroup setting for conmon
conmon_cgroup = "{{ .ConmonCgroup }}"

# Environment variable list for the conmon process, used for passing necessary
# environment variables to conmon or the runtime.
conmon_env = [
{{ range $env := .ConmonEnv }}{{ printf "\t%q,\n" $env }}{{ end }}]

# If true, SELinux will be used for pod separation on the host.
selinux = {{ .SELinux }}

# Path to the seccomp.json profile which is used as the default seccomp profile
# for the runtime. If not specified, then the internal default seccomp profile
# will be used.
seccomp_profile = "{{ .SeccompProfile }}"

# Used to change the name of the default AppArmor profile of CRI-O. The default
# profile name is "crio-default-" followed by the version string of CRI-O.
apparmor_profile = "{{ .ApparmorProfile }}"

# Cgroup management implementation used for the runtime.
cgroup_manager = "{{ .CgroupManager }}"

# List of default capabilities for containers. If it is empty or commented out,
# only the capabilities defined in the containers json file by the user/kube
# will be added.
default_capabilities = [
{{ range $capability := .DefaultCapabilities}}{{ printf "\t%q, \n" $capability}}{{ end }}]

# List of default sysctls. If it is empty or commented out, only the sysctls
# defined in the container json file by the user/kube will be added.
default_sysctls = [
{{ range $sysctl := .DefaultSysctls}}{{ printf "\t%q, \n" $sysctl}}{{ end }}]

# List of additional devices. specified as
# "<device-on-host>:<device-on-container>:<permissions>", for example: "--device=/dev/sdc:/dev/xvdc:rwm".
#If it is empty or commented out, only the devices
# defined in the container json file by the user/kube will be added.
additional_devices = [
{{ range $device := .AdditionalDevices}}{{ printf "\t%q, \n" $device}}{{ end }}]

# Path to OCI hooks directories for automatically executed hooks.
hooks_dir = [
{{ range $hooksDir := .HooksDir }}{{ printf "\t%q, \n" $hooksDir}}{{ end }}]

# List of default mounts for each container. **Deprecated:** this option will
# be removed in future versions in favor of default_mounts_file.
default_mounts = [
{{ range $mount := .DefaultMounts }}{{ printf "\t%q, \n" $mount }}{{ end }}]

# Path to the file specifying the defaults mounts for each container. The
# format of the config is /SRC:/DST, one mount per line. Notice that CRI-O reads
# its default mounts from the following two files:
#
#   1) /etc/containers/mounts.conf (i.e., default_mounts_file): This is the
#      override file, where users can either add in their own default mounts, or
#      override the default mounts shipped with the package.
#
#   2) /usr/share/containers/mounts.conf: This is the default file read for
#      mounts. If you want CRI-O to read from a different, specific mounts file,
#      you can change the default_mounts_file. Note, if this is done, CRI-O will
#      only add mounts it finds in this file.
#
#default_mounts_file = "{{ .DefaultMountsFile }}"

# Maximum number of processes allowed in a container.
pids_limit = {{ .PidsLimit }}

# Maximum sized allowed for the container log file. Negative numbers indicate
# that no size limit is imposed. If it is positive, it must be >= 8192 to
# match/exceed conmon's read buffer. The file is truncated and re-opened so the
# limit is never exceeded.
log_size_max = {{ .LogSizeMax }}

# Whether container output should be logged to journald in addition to the kuberentes log file
log_to_journald = {{ .LogToJournald }}

# Path to directory in which container exit files are written to by conmon.
container_exits_dir = "{{ .ContainerExitsDir }}"

# Path to directory for container attach sockets.
container_attach_socket_dir = "{{ .ContainerAttachSocketDir }}"

# If set to true, all containers will run in read-only mode.
read_only = {{ .ReadOnly }}

# Changes the verbosity of the logs based on the level it is set to. Options
# are fatal, panic, error, warn, info, and debug. This option supports live
# configuration reload.
log_level = "{{ .LogLevel }}"

# The default log directory where all logs will go unless directly specified by the kubelet
log_dir = "{{ .LogDir }}"

# The UID mappings for the user namespace of each container. A range is
# specified in the form containerUID:HostUID:Size. Multiple ranges must be
# separated by comma.
uid_mappings = "{{ .UIDMappings }}"

# The GID mappings for the user namespace of each container. A range is
# specified in the form containerGID:HostGID:Size. Multiple ranges must be
# separated by comma.
gid_mappings = "{{ .GIDMappings }}"

# The minimal amount of time in seconds to wait before issuing a timeout
# regarding the proper termination of the container.
ctr_stop_timeout = {{ .CtrStopTimeout }}

# ManageNetworkNSLifecycle determines whether we pin and remove network namespace
# and manage its lifecycle.
manage_network_ns_lifecycle = {{ .ManageNetworkNSLifecycle }}

# The "crio.runtime.runtimes" table defines a list of OCI compatible runtimes.
# The runtime to use is picked based on the runtime_handler provided by the CRI.
# If no runtime_handler is provided, the runtime will be picked based on the level
# of trust of the workload.
{{ range $runtime_name, $runtime_handler := .Runtimes  }}
[crio.runtime.runtimes.{{ $runtime_name }}]
runtime_path = "{{ $runtime_handler.RuntimePath }}"
runtime_type = "{{ $runtime_handler.RuntimeType }}"
runtime_root = "{{ $runtime_handler.RuntimeRoot }}"
{{ end }}

# The crio.image table contains settings pertaining to the management of OCI images.
#
# CRI-O reads its configured registries defaults from the system wide
# containers-registries.conf(5) located in /etc/containers/registries.conf. If
# you want to modify just CRI-O, you can change the registries configuration in
# this file. Otherwise, leave insecure_registries and registries commented out to
# use the system's defaults from /etc/containers/registries.conf.
[crio.image]

# Default transport for pulling images from a remote container storage.
default_transport = "{{ .DefaultTransport }}"

# The path to a file containing credentials necessary for pulling images from
# secure registries. The file is similar to that of /var/lib/kubelet/config.json
global_auth_file = "{{ .GlobalAuthFile }}"

# The image used to instantiate infra containers.
# This option supports live configuration reload.
pause_image = "{{ .PauseImage }}"

# The path to a file containing credentials specific for pulling the pause_image from
# above. The file is similar to that of /var/lib/kubelet/config.json
# This option supports live configuration reload.
pause_image_auth_file = "{{ .PauseImageAuthFile }}"

# The command to run to have a container stay in the paused state.
# This option supports live configuration reload.
pause_command = "{{ .PauseCommand }}"

# Path to the file which decides what sort of policy we use when deciding
# whether or not to trust an image that we've pulled. It is not recommended that
# this option be used, as the default behavior of using the system-wide default
# policy (i.e., /etc/containers/policy.json) is most often preferred. Please
# refer to containers-policy.json(5) for more details.
signature_policy = "{{ .SignaturePolicyPath }}"

# Controls how image volumes are handled. The valid values are mkdir, bind and
# ignore; the latter will ignore volumes entirely.
image_volumes = "{{ .ImageVolumes }}"

# List of registries to be used when pulling an unqualified image (e.g.,
# "alpine:latest"). By default, registries is set to "docker.io" for
# compatibility reasons. Depending on your workload and usecase you may add more
# registries (e.g., "quay.io", "registry.fedoraproject.org",
# "registry.opensuse.org", etc.).
#registries = [
# {{ range $opt := .Registries }}{{ printf "\t%q,\n#" $opt }}{{ end }}]

# The crio.network table containers settings pertaining to the management of
# CNI plugins.
[crio.network]

# Path to the directory where CNI configuration files are located.
network_dir = "{{ .NetworkDir }}"

# Paths to directories where CNI plugin binaries are located.
plugin_dirs = [
{{ range $opt := .PluginDirs }}{{ printf "\t%q,\n" $opt }}{{ end }}]
`))

// TODO: Currently ImageDir isn't really used, so we haven't added it to this
//       template. Add it once the storage code has been merged.

// NOTE: please propagate any changes to the template to docs/crio.conf.5.md

var configCommand = cli.Command{
	Name:  "config",
	Usage: "generate crio configuration files",
	Flags: []cli.Flag{
		cli.BoolFlag{
			Name:  "default",
			Usage: "output the default configuration",
		},
	},
	Action: func(c *cli.Context) error {
		var err error
		// At this point, app.Before has already parsed the user's chosen
		// config file. So no need to handle that here.
		config := c.App.Metadata["config"].(*server.Config) // nolint: errcheck
		systemContext := &types.SystemContext{}
		if c.Bool("default") {
			config, err = server.DefaultConfig()
			if err != nil {
				return err
			}
		}

		// Validate the configuration during generation
		if err = config.Validate(systemContext, false); err != nil {
			return err
		}

		// Output the commented config.
		return commentedConfigTemplate.ExecuteTemplate(os.Stdout, "config", config)
	},
}
<commit_msg>Add missing `insecure_registries` to config template<commit_after>package main

import (
	"os"
	"text/template"

	"github.com/containers/image/types"
	"github.com/cri-o/cri-o/server"
	"github.com/urfave/cli"
)

// NOTE: please propagate any changes to the template to docs/crio.conf.5.md

var commentedConfigTemplate = template.Must(template.New("config").Parse(`
# The CRI-O configuration file specifies all of the available configuration
# options and command-line flags for the crio(8) OCI Kubernetes Container Runtime
# daemon, but in a TOML format that can be more easily modified and versioned.
#
# Please refer to crio.conf(5) for details of all configuration options.

# CRI-O supports partial configuration reload during runtime, which can be
# done by sending SIGHUP to the running process. Currently supported options
# are explicitly mentioned with: 'This option supports live configuration
# reload'.

# CRI-O reads its storage defaults from the containers-storage.conf(5) file
# located at /etc/containers/storage.conf. Modify this storage configuration if
# you want to change the system's defaults. If you want to modify storage just
# for CRI-O, you can change the storage configuration options here.
[crio]

# Path to the "root directory". CRI-O stores all of its data, including
# containers images, in this directory.
#root = "{{ .Root }}"

# Path to the "run directory". CRI-O stores all of its state in this directory.
#runroot = "{{ .RunRoot }}"

# Storage driver used to manage the storage of images and containers. Please
# refer to containers-storage.conf(5) to see all available storage drivers.
#storage_driver = "{{ .Storage }}"

# List to pass options to the storage driver. Please refer to
# containers-storage.conf(5) to see all available storage options.
#storage_option = [
{{ range $opt := .StorageOptions }}{{ printf "#\t%q,\n" $opt }}{{ end }}#]

# If set to false, in-memory locking will be used instead of file-based locking.
# **Deprecated** this option will be removed in the future.
file_locking = {{ .FileLocking }}

# Path to the lock file.
# **Deprecated** this option will be removed in the future.
file_locking_path = "{{ .FileLockingPath }}"

# The crio.api table contains settings for the kubelet/gRPC interface.
[crio.api]

# Path to AF_LOCAL socket on which CRI-O will listen.
listen = "{{ .Listen }}"

# IP address on which the stream server will listen.
stream_address = "{{ .StreamAddress }}"

# The port on which the stream server will listen.
stream_port = "{{ .StreamPort }}"

# Enable encrypted TLS transport of the stream server.
stream_enable_tls = {{ .StreamEnableTLS }}

# Path to the x509 certificate file used to serve the encrypted stream. This
# file can change, and CRI-O will automatically pick up the changes within 5
# minutes.
stream_tls_cert = "{{ .StreamTLSCert }}"

# Path to the key file used to serve the encrypted stream. This file can
# change, and CRI-O will automatically pick up the changes within 5 minutes.
stream_tls_key = "{{ .StreamTLSKey }}"

# Path to the x509 CA(s) file used to verify and authenticate client
# communication with the encrypted stream. This file can change, and CRI-O will
# automatically pick up the changes within 5 minutes.
stream_tls_ca = "{{ .StreamTLSCA }}"

# Maximum grpc send message size in bytes. If not set or <=0, then CRI-O will default to 16 * 1024 * 1024.
grpc_max_send_msg_size = {{ .GRPCMaxSendMsgSize }}

# Maximum grpc receive message size. If not set or <= 0, then CRI-O will default to 16 * 1024 * 1024.
grpc_max_recv_msg_size = {{ .GRPCMaxRecvMsgSize }}

# The crio.runtime table contains settings pertaining to the OCI runtime used
# and options for how to set up and manage the OCI runtime.
[crio.runtime]

# A list of ulimits to be set in containers by default, specified as
# "<ulimit name>=<soft limit>:<hard limit>", for example:
# "nofile=1024:2048"
# If nothing is set here, settings will be inherited from the CRI-O daemon
#default_ulimits = [
{{ range $ulimit := .DefaultUlimits }}{{ printf "#\t%q,\n" $ulimit }}{{ end }}#]

# default_runtime is the _name_ of the OCI runtime to be used as the default.
# The name is matched against the runtimes map below.
default_runtime = "{{ .DefaultRuntime }}"

# If true, the runtime will not use pivot_root, but instead use MS_MOVE.
no_pivot = {{ .NoPivot }}

# Path to the conmon binary, used for monitoring the OCI runtime.
# Will be searched for using $PATH if empty.
conmon = "{{ .Conmon }}"

# Cgroup setting for conmon
conmon_cgroup = "{{ .ConmonCgroup }}"

# Environment variable list for the conmon process, used for passing necessary
# environment variables to conmon or the runtime.
conmon_env = [
{{ range $env := .ConmonEnv }}{{ printf "\t%q,\n" $env }}{{ end }}]

# If true, SELinux will be used for pod separation on the host.
selinux = {{ .SELinux }}

# Path to the seccomp.json profile which is used as the default seccomp profile
# for the runtime. If not specified, then the internal default seccomp profile
# will be used.
seccomp_profile = "{{ .SeccompProfile }}"

# Used to change the name of the default AppArmor profile of CRI-O. The default
# profile name is "crio-default-" followed by the version string of CRI-O.
apparmor_profile = "{{ .ApparmorProfile }}"

# Cgroup management implementation used for the runtime.
cgroup_manager = "{{ .CgroupManager }}"

# List of default capabilities for containers. If it is empty or commented out,
# only the capabilities defined in the containers json file by the user/kube
# will be added.
default_capabilities = [
{{ range $capability := .DefaultCapabilities}}{{ printf "\t%q, \n" $capability}}{{ end }}]

# List of default sysctls. If it is empty or commented out, only the sysctls
# defined in the container json file by the user/kube will be added.
default_sysctls = [
{{ range $sysctl := .DefaultSysctls}}{{ printf "\t%q, \n" $sysctl}}{{ end }}]

# List of additional devices. specified as
# "<device-on-host>:<device-on-container>:<permissions>", for example: "--device=/dev/sdc:/dev/xvdc:rwm".
#If it is empty or commented out, only the devices
# defined in the container json file by the user/kube will be added.
additional_devices = [
{{ range $device := .AdditionalDevices}}{{ printf "\t%q, \n" $device}}{{ end }}]

# Path to OCI hooks directories for automatically executed hooks.
hooks_dir = [
{{ range $hooksDir := .HooksDir }}{{ printf "\t%q, \n" $hooksDir}}{{ end }}]

# List of default mounts for each container. **Deprecated:** this option will
# be removed in future versions in favor of default_mounts_file.
default_mounts = [
{{ range $mount := .DefaultMounts }}{{ printf "\t%q, \n" $mount }}{{ end }}]

# Path to the file specifying the defaults mounts for each container. The
# format of the config is /SRC:/DST, one mount per line. Notice that CRI-O reads
# its default mounts from the following two files:
#
#   1) /etc/containers/mounts.conf (i.e., default_mounts_file): This is the
#      override file, where users can either add in their own default mounts, or
#      override the default mounts shipped with the package.
#
#   2) /usr/share/containers/mounts.conf: This is the default file read for
#      mounts. If you want CRI-O to read from a different, specific mounts file,
#      you can change the default_mounts_file. Note, if this is done, CRI-O will
#      only add mounts it finds in this file.
#
#default_mounts_file = "{{ .DefaultMountsFile }}"

# Maximum number of processes allowed in a container.
pids_limit = {{ .PidsLimit }}

# Maximum sized allowed for the container log file. Negative numbers indicate
# that no size limit is imposed. If it is positive, it must be >= 8192 to
# match/exceed conmon's read buffer. The file is truncated and re-opened so the
# limit is never exceeded.
log_size_max = {{ .LogSizeMax }}

# Whether container output should be logged to journald in addition to the kuberentes log file
log_to_journald = {{ .LogToJournald }}

# Path to directory in which container exit files are written to by conmon.
container_exits_dir = "{{ .ContainerExitsDir }}"

# Path to directory for container attach sockets.
container_attach_socket_dir = "{{ .ContainerAttachSocketDir }}"

# If set to true, all containers will run in read-only mode.
read_only = {{ .ReadOnly }}

# Changes the verbosity of the logs based on the level it is set to. Options
# are fatal, panic, error, warn, info, and debug. This option supports live
# configuration reload.
log_level = "{{ .LogLevel }}"

# The default log directory where all logs will go unless directly specified by the kubelet
log_dir = "{{ .LogDir }}"

# The UID mappings for the user namespace of each container. A range is
# specified in the form containerUID:HostUID:Size. Multiple ranges must be
# separated by comma.
uid_mappings = "{{ .UIDMappings }}"

# The GID mappings for the user namespace of each container. A range is
# specified in the form containerGID:HostGID:Size. Multiple ranges must be
# separated by comma.
gid_mappings = "{{ .GIDMappings }}"

# The minimal amount of time in seconds to wait before issuing a timeout
# regarding the proper termination of the container.
ctr_stop_timeout = {{ .CtrStopTimeout }}

# ManageNetworkNSLifecycle determines whether we pin and remove network namespace
# and manage its lifecycle.
manage_network_ns_lifecycle = {{ .ManageNetworkNSLifecycle }}

# The "crio.runtime.runtimes" table defines a list of OCI compatible runtimes.
# The runtime to use is picked based on the runtime_handler provided by the CRI.
# If no runtime_handler is provided, the runtime will be picked based on the level
# of trust of the workload.
{{ range $runtime_name, $runtime_handler := .Runtimes  }}
[crio.runtime.runtimes.{{ $runtime_name }}]
runtime_path = "{{ $runtime_handler.RuntimePath }}"
runtime_type = "{{ $runtime_handler.RuntimeType }}"
runtime_root = "{{ $runtime_handler.RuntimeRoot }}"
{{ end }}

# The crio.image table contains settings pertaining to the management of OCI images.
#
# CRI-O reads its configured registries defaults from the system wide
# containers-registries.conf(5) located in /etc/containers/registries.conf. If
# you want to modify just CRI-O, you can change the registries configuration in
# this file. Otherwise, leave insecure_registries and registries commented out to
# use the system's defaults from /etc/containers/registries.conf.
[crio.image]

# Default transport for pulling images from a remote container storage.
default_transport = "{{ .DefaultTransport }}"

# The path to a file containing credentials necessary for pulling images from
# secure registries. The file is similar to that of /var/lib/kubelet/config.json
global_auth_file = "{{ .GlobalAuthFile }}"

# The image used to instantiate infra containers.
# This option supports live configuration reload.
pause_image = "{{ .PauseImage }}"

# The path to a file containing credentials specific for pulling the pause_image from
# above. The file is similar to that of /var/lib/kubelet/config.json
# This option supports live configuration reload.
pause_image_auth_file = "{{ .PauseImageAuthFile }}"

# The command to run to have a container stay in the paused state.
# This option supports live configuration reload.
pause_command = "{{ .PauseCommand }}"

# Path to the file which decides what sort of policy we use when deciding
# whether or not to trust an image that we've pulled. It is not recommended that
# this option be used, as the default behavior of using the system-wide default
# policy (i.e., /etc/containers/policy.json) is most often preferred. Please
# refer to containers-policy.json(5) for more details.
signature_policy = "{{ .SignaturePolicyPath }}"

# List of registries to skip TLS verification for pulling images. Please
# consider configuring the registries via /etc/containers/registries.conf before
# changing them here.
#insecure_registries = "{{ .InsecureRegistries }}"

# Controls how image volumes are handled. The valid values are mkdir, bind and
# ignore; the latter will ignore volumes entirely.
image_volumes = "{{ .ImageVolumes }}"

# List of registries to be used when pulling an unqualified image (e.g.,
# "alpine:latest"). By default, registries is set to "docker.io" for
# compatibility reasons. Depending on your workload and usecase you may add more
# registries (e.g., "quay.io", "registry.fedoraproject.org",
# "registry.opensuse.org", etc.).
#registries = [
# {{ range $opt := .Registries }}{{ printf "\t%q,\n#" $opt }}{{ end }}]

# The crio.network table containers settings pertaining to the management of
# CNI plugins.
[crio.network]

# Path to the directory where CNI configuration files are located.
network_dir = "{{ .NetworkDir }}"

# Paths to directories where CNI plugin binaries are located.
plugin_dirs = [
{{ range $opt := .PluginDirs }}{{ printf "\t%q,\n" $opt }}{{ end }}]
`))

// TODO: Currently ImageDir isn't really used, so we haven't added it to this
//       template. Add it once the storage code has been merged.

// NOTE: please propagate any changes to the template to docs/crio.conf.5.md

var configCommand = cli.Command{
	Name:  "config",
	Usage: "generate crio configuration files",
	Flags: []cli.Flag{
		cli.BoolFlag{
			Name:  "default",
			Usage: "output the default configuration",
		},
	},
	Action: func(c *cli.Context) error {
		var err error
		// At this point, app.Before has already parsed the user's chosen
		// config file. So no need to handle that here.
		config := c.App.Metadata["config"].(*server.Config) // nolint: errcheck
		systemContext := &types.SystemContext{}
		if c.Bool("default") {
			config, err = server.DefaultConfig()
			if err != nil {
				return err
			}
		}

		// Validate the configuration during generation
		if err = config.Validate(systemContext, false); err != nil {
			return err
		}

		// Output the commented config.
		return commentedConfigTemplate.ExecuteTemplate(os.Stdout, "config", config)
	},
}
<|endoftext|>


<commit_before>package cmd

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/AlecAivazis/survey/v2"
	"github.com/AlecAivazis/survey/v2/terminal"
	"github.com/apex/log"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/pkg/parsers/kernel"
	"github.com/docker/docker/pkg/parsers/operatingsystem"
	"github.com/spf13/cobra"

	"github.com/pterodactyl/wings/config"
	"github.com/pterodactyl/wings/environment"
	"github.com/pterodactyl/wings/loggers/cli"
	"github.com/pterodactyl/wings/system"
)

const (
	DefaultHastebinUrl = "https://ptero.co"
	DefaultLogLines    = 200
)

var diagnosticsArgs struct {
	IncludeEndpoints   bool
	IncludeLogs        bool
	ReviewBeforeUpload bool
	HastebinURL        string
	LogLines           int
}

func newDiagnosticsCommand() *cobra.Command {
	command := &cobra.Command{
		Use:   "diagnostics",
		Short: "Collect and report information about this Wings instance to assist in debugging.",
		PreRun: func(cmd *cobra.Command, args []string) {
			initConfig()
			log.SetHandler(cli.Default)
		},
		Run: diagnosticsCmdRun,
	}

	command.Flags().StringVar(&diagnosticsArgs.HastebinURL, "hastebin-url", DefaultHastebinUrl, "the url of the hastebin instance to use")
	command.Flags().IntVar(&diagnosticsArgs.LogLines, "log-lines", DefaultLogLines, "the number of log lines to include in the report")

	return command
}

// diagnosticsCmdRun collects diagnostics about wings, it's configuration and the node.
// We collect:
// - wings and docker versions
// - relevant parts of daemon configuration
// - the docker debug output
// - running docker containers
// - logs
func diagnosticsCmdRun(cmd *cobra.Command, args []string) {
	questions := []*survey.Question{
		{
			Name:   "IncludeEndpoints",
			Prompt: &survey.Confirm{Message: "Do you want to include endpoints (i.e. the FQDN/IP of your panel)?", Default: false},
		},
		{
			Name:   "IncludeLogs",
			Prompt: &survey.Confirm{Message: "Do you want to include the latest logs?", Default: true},
		},
		{
			Name: "ReviewBeforeUpload",
			Prompt: &survey.Confirm{
				Message: "Do you want to review the collected data before uploading to " + diagnosticsArgs.HastebinURL + "?",
				Help:    "The data, especially the logs, might contain sensitive information, so you should review it. You will be asked again if you want to upload.",
				Default: true,
			},
		},
	}
	if err := survey.Ask(questions, &diagnosticsArgs); err != nil {
		if err == terminal.InterruptErr {
			return
		}
		panic(err)
	}

	dockerVersion, dockerInfo, dockerErr := getDockerInfo()

	output := &strings.Builder{}
	fmt.Fprintln(output, "Pterodactyl Wings - Diagnostics Report")
	printHeader(output, "Versions")
	fmt.Fprintln(output, "               Wings:", system.Version)
	if dockerErr == nil {
		fmt.Fprintln(output, "              Docker:", dockerVersion.Version)
	}
	if v, err := kernel.GetKernelVersion(); err == nil {
		fmt.Fprintln(output, "              Kernel:", v)
	}
	if os, err := operatingsystem.GetOperatingSystem(); err == nil {
		fmt.Fprintln(output, "                  OS:", os)
	}

	printHeader(output, "Wings Configuration")
	if err := config.FromFile(config.DefaultLocation); err != nil {
	}
	cfg := config.Get()
	fmt.Fprintln(output, "      Panel Location:", redact(cfg.PanelLocation))
	fmt.Fprintln(output, "")
	fmt.Fprintln(output, "  Internal Webserver:", redact(cfg.Api.Host), ":", cfg.Api.Port)
	fmt.Fprintln(output, "         SSL Enabled:", cfg.Api.Ssl.Enabled)
	fmt.Fprintln(output, "     SSL Certificate:", redact(cfg.Api.Ssl.CertificateFile))
	fmt.Fprintln(output, "             SSL Key:", redact(cfg.Api.Ssl.KeyFile))
	fmt.Fprintln(output, "")
	fmt.Fprintln(output, "         SFTP Server:", redact(cfg.System.Sftp.Address), ":", cfg.System.Sftp.Port)
	fmt.Fprintln(output, "      SFTP Read-Only:", cfg.System.Sftp.ReadOnly)
	fmt.Fprintln(output, "")
	fmt.Fprintln(output, "      Root Directory:", cfg.System.RootDirectory)
	fmt.Fprintln(output, "      Logs Directory:", cfg.System.LogDirectory)
	fmt.Fprintln(output, "      Data Directory:", cfg.System.Data)
	fmt.Fprintln(output, "   Archive Directory:", cfg.System.ArchiveDirectory)
	fmt.Fprintln(output, "    Backup Directory:", cfg.System.BackupDirectory)
	fmt.Fprintln(output, "")
	fmt.Fprintln(output, "            Username:", cfg.System.Username)
	fmt.Fprintln(output, "         Server Time:", time.Now().Format(time.RFC1123Z))
	fmt.Fprintln(output, "          Debug Mode:", cfg.Debug)

	printHeader(output, "Docker: Info")
	if dockerErr == nil {
		fmt.Fprintln(output, "Server Version:", dockerInfo.ServerVersion)
		fmt.Fprintln(output, "Storage Driver:", dockerInfo.Driver)
		if dockerInfo.DriverStatus != nil {
			for _, pair := range dockerInfo.DriverStatus {
				fmt.Fprintf(output, "  %s: %s\n", pair[0], pair[1])
			}
		}
		if dockerInfo.SystemStatus != nil {
			for _, pair := range dockerInfo.SystemStatus {
				fmt.Fprintf(output, " %s: %s\n", pair[0], pair[1])
			}
		}
		fmt.Fprintln(output, "LoggingDriver:", dockerInfo.LoggingDriver)
		fmt.Fprintln(output, " CgroupDriver:", dockerInfo.CgroupDriver)
		if len(dockerInfo.Warnings) > 0 {
			for _, w := range dockerInfo.Warnings {
				fmt.Fprintln(output, w)
			}
		}
	} else {
		fmt.Fprintln(output, dockerErr.Error())
	}

	printHeader(output, "Docker: Running Containers")
	c := exec.Command("docker", "ps")
	if co, err := c.Output(); err == nil {
		output.Write(co)
	} else {
		fmt.Fprint(output, "Couldn't list containers: ", err)
	}

	printHeader(output, "Latest Wings Logs")
	if diagnosticsArgs.IncludeLogs {
		p := "/var/log/pterodactyl/wings.log"
		if cfg != nil {
			p = path.Join(cfg.System.LogDirectory, "wings.log")
		}
		if c, err := exec.Command("tail", "-n", strconv.Itoa(diagnosticsArgs.LogLines), p).Output(); err != nil {
			fmt.Fprintln(output, "No logs found or an error occurred.")
		} else {
			fmt.Fprintf(output, "%s\n", string(c))
		}
	} else {
		fmt.Fprintln(output, "Logs redacted.")
	}

	fmt.Println("\n---------------  generated report  ---------------")
	fmt.Println(output.String())
	fmt.Print("---------------   end of report    ---------------\n\n")

	upload := !diagnosticsArgs.ReviewBeforeUpload
	if !upload {
		survey.AskOne(&survey.Confirm{Message: "Upload to " + diagnosticsArgs.HastebinURL + "?", Default: false}, &upload)
	}
	if upload {
		if !diagnosticsArgs.IncludeEndpoints {
			s := output.String()
			output.Reset()
			a := strings.ReplaceAll(cfg.PanelLocation, s, "{redacted}")
			a = strings.ReplaceAll(cfg.Api.Host, a, "{redacted}")
			a = strings.ReplaceAll(cfg.Api.Ssl.CertificateFile, a, "{redacted}")
			a = strings.ReplaceAll(cfg.Api.Ssl.KeyFile, a, "{redacted}")
			a = strings.ReplaceAll(cfg.System.Sftp.Address, a, "{redacted}")
			output.WriteString(a)
		}
		u, err := uploadToHastebin(diagnosticsArgs.HastebinURL, output.String())
		if err == nil {
			fmt.Println("Your report is available here: ", u)
		}
	}
}

func getDockerInfo() (types.Version, types.Info, error) {
	client, err := environment.Docker()
	if err != nil {
		return types.Version{}, types.Info{}, err
	}
	dockerVersion, err := client.ServerVersion(context.Background())
	if err != nil {
		return types.Version{}, types.Info{}, err
	}
	dockerInfo, err := client.Info(context.Background())
	if err != nil {
		return types.Version{}, types.Info{}, err
	}
	return dockerVersion, dockerInfo, nil
}

func uploadToHastebin(hbUrl, content string) (string, error) {
	r := strings.NewReader(content)
	u, err := url.Parse(hbUrl)
	if err != nil {
		return "", err
	}
	u.Path = path.Join(u.Path, "documents")
	res, err := http.Post(u.String(), "plain/text", r)
	if err != nil || res.StatusCode != 200 {
		fmt.Println("Failed to upload report to ", u.String(), err)
		return "", err
	}
	pres := make(map[string]interface{})
	body, err := io.ReadAll(res.Body)
	if err != nil {
		fmt.Println("Failed to parse response.", err)
		return "", err
	}
	json.Unmarshal(body, &pres)
	if key, ok := pres["key"].(string); ok {
		u, _ := url.Parse(hbUrl)
		u.Path = path.Join(u.Path, key)
		return u.String(), nil
	}
	return "", errors.New("failed to find key in response")
}

func redact(s string) string {
	if !diagnosticsArgs.IncludeEndpoints {
		return "{redacted}"
	}
	return s
}

func printHeader(w io.Writer, title string) {
	fmt.Fprintln(w, "\n|\n|", title)
	fmt.Fprintln(w, "| ------------------------------")
}
<commit_msg>diagnostics: properly redact endpoints<commit_after>package cmd

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/AlecAivazis/survey/v2"
	"github.com/AlecAivazis/survey/v2/terminal"
	"github.com/apex/log"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/pkg/parsers/kernel"
	"github.com/docker/docker/pkg/parsers/operatingsystem"
	"github.com/spf13/cobra"

	"github.com/pterodactyl/wings/config"
	"github.com/pterodactyl/wings/environment"
	"github.com/pterodactyl/wings/loggers/cli"
	"github.com/pterodactyl/wings/system"
)

const (
	DefaultHastebinUrl = "https://ptero.co"
	DefaultLogLines    = 200
)

var diagnosticsArgs struct {
	IncludeEndpoints   bool
	IncludeLogs        bool
	ReviewBeforeUpload bool
	HastebinURL        string
	LogLines           int
}

func newDiagnosticsCommand() *cobra.Command {
	command := &cobra.Command{
		Use:   "diagnostics",
		Short: "Collect and report information about this Wings instance to assist in debugging.",
		PreRun: func(cmd *cobra.Command, args []string) {
			initConfig()
			log.SetHandler(cli.Default)
		},
		Run: diagnosticsCmdRun,
	}

	command.Flags().StringVar(&diagnosticsArgs.HastebinURL, "hastebin-url", DefaultHastebinUrl, "the url of the hastebin instance to use")
	command.Flags().IntVar(&diagnosticsArgs.LogLines, "log-lines", DefaultLogLines, "the number of log lines to include in the report")

	return command
}

// diagnosticsCmdRun collects diagnostics about wings, it's configuration and the node.
// We collect:
// - wings and docker versions
// - relevant parts of daemon configuration
// - the docker debug output
// - running docker containers
// - logs
func diagnosticsCmdRun(cmd *cobra.Command, args []string) {
	questions := []*survey.Question{
		{
			Name:   "IncludeEndpoints",
			Prompt: &survey.Confirm{Message: "Do you want to include endpoints (i.e. the FQDN/IP of your panel)?", Default: false},
		},
		{
			Name:   "IncludeLogs",
			Prompt: &survey.Confirm{Message: "Do you want to include the latest logs?", Default: true},
		},
		{
			Name: "ReviewBeforeUpload",
			Prompt: &survey.Confirm{
				Message: "Do you want to review the collected data before uploading to " + diagnosticsArgs.HastebinURL + "?",
				Help:    "The data, especially the logs, might contain sensitive information, so you should review it. You will be asked again if you want to upload.",
				Default: true,
			},
		},
	}
	if err := survey.Ask(questions, &diagnosticsArgs); err != nil {
		if err == terminal.InterruptErr {
			return
		}
		panic(err)
	}

	dockerVersion, dockerInfo, dockerErr := getDockerInfo()

	output := &strings.Builder{}
	fmt.Fprintln(output, "Pterodactyl Wings - Diagnostics Report")
	printHeader(output, "Versions")
	fmt.Fprintln(output, "               Wings:", system.Version)
	if dockerErr == nil {
		fmt.Fprintln(output, "              Docker:", dockerVersion.Version)
	}
	if v, err := kernel.GetKernelVersion(); err == nil {
		fmt.Fprintln(output, "              Kernel:", v)
	}
	if os, err := operatingsystem.GetOperatingSystem(); err == nil {
		fmt.Fprintln(output, "                  OS:", os)
	}

	printHeader(output, "Wings Configuration")
	if err := config.FromFile(config.DefaultLocation); err != nil {
	}
	cfg := config.Get()
	fmt.Fprintln(output, "      Panel Location:", redact(cfg.PanelLocation))
	fmt.Fprintln(output, "")
	fmt.Fprintln(output, "  Internal Webserver:", redact(cfg.Api.Host), ":", cfg.Api.Port)
	fmt.Fprintln(output, "         SSL Enabled:", cfg.Api.Ssl.Enabled)
	fmt.Fprintln(output, "     SSL Certificate:", redact(cfg.Api.Ssl.CertificateFile))
	fmt.Fprintln(output, "             SSL Key:", redact(cfg.Api.Ssl.KeyFile))
	fmt.Fprintln(output, "")
	fmt.Fprintln(output, "         SFTP Server:", redact(cfg.System.Sftp.Address), ":", cfg.System.Sftp.Port)
	fmt.Fprintln(output, "      SFTP Read-Only:", cfg.System.Sftp.ReadOnly)
	fmt.Fprintln(output, "")
	fmt.Fprintln(output, "      Root Directory:", cfg.System.RootDirectory)
	fmt.Fprintln(output, "      Logs Directory:", cfg.System.LogDirectory)
	fmt.Fprintln(output, "      Data Directory:", cfg.System.Data)
	fmt.Fprintln(output, "   Archive Directory:", cfg.System.ArchiveDirectory)
	fmt.Fprintln(output, "    Backup Directory:", cfg.System.BackupDirectory)
	fmt.Fprintln(output, "")
	fmt.Fprintln(output, "            Username:", cfg.System.Username)
	fmt.Fprintln(output, "         Server Time:", time.Now().Format(time.RFC1123Z))
	fmt.Fprintln(output, "          Debug Mode:", cfg.Debug)

	printHeader(output, "Docker: Info")
	if dockerErr == nil {
		fmt.Fprintln(output, "Server Version:", dockerInfo.ServerVersion)
		fmt.Fprintln(output, "Storage Driver:", dockerInfo.Driver)
		if dockerInfo.DriverStatus != nil {
			for _, pair := range dockerInfo.DriverStatus {
				fmt.Fprintf(output, "  %s: %s\n", pair[0], pair[1])
			}
		}
		if dockerInfo.SystemStatus != nil {
			for _, pair := range dockerInfo.SystemStatus {
				fmt.Fprintf(output, " %s: %s\n", pair[0], pair[1])
			}
		}
		fmt.Fprintln(output, "LoggingDriver:", dockerInfo.LoggingDriver)
		fmt.Fprintln(output, " CgroupDriver:", dockerInfo.CgroupDriver)
		if len(dockerInfo.Warnings) > 0 {
			for _, w := range dockerInfo.Warnings {
				fmt.Fprintln(output, w)
			}
		}
	} else {
		fmt.Fprintln(output, dockerErr.Error())
	}

	printHeader(output, "Docker: Running Containers")
	c := exec.Command("docker", "ps")
	if co, err := c.Output(); err == nil {
		output.Write(co)
	} else {
		fmt.Fprint(output, "Couldn't list containers: ", err)
	}

	printHeader(output, "Latest Wings Logs")
	if diagnosticsArgs.IncludeLogs {
		p := "/var/log/pterodactyl/wings.log"
		if cfg != nil {
			p = path.Join(cfg.System.LogDirectory, "wings.log")
		}
		if c, err := exec.Command("tail", "-n", strconv.Itoa(diagnosticsArgs.LogLines), p).Output(); err != nil {
			fmt.Fprintln(output, "No logs found or an error occurred.")
		} else {
			fmt.Fprintf(output, "%s\n", string(c))
		}
	} else {
		fmt.Fprintln(output, "Logs redacted.")
	}

	if !diagnosticsArgs.IncludeEndpoints {
		s := output.String()
		output.Reset()
		s = strings.ReplaceAll(s, cfg.PanelLocation, "{redacted}")
		s = strings.ReplaceAll(s, cfg.Api.Host, "{redacted}")
		s = strings.ReplaceAll(s, cfg.Api.Ssl.CertificateFile, "{redacted}")
		s = strings.ReplaceAll(s, cfg.Api.Ssl.KeyFile, "{redacted}")
		s = strings.ReplaceAll(s, cfg.System.Sftp.Address, "{redacted}")
		output.WriteString(s)
	}

	fmt.Println("\n---------------  generated report  ---------------")
	fmt.Println(output.String())
	fmt.Print("---------------   end of report    ---------------\n\n")

	upload := !diagnosticsArgs.ReviewBeforeUpload
	if !upload {
		survey.AskOne(&survey.Confirm{Message: "Upload to " + diagnosticsArgs.HastebinURL + "?", Default: false}, &upload)
	}
	if upload {
		u, err := uploadToHastebin(diagnosticsArgs.HastebinURL, output.String())
		if err == nil {
			fmt.Println("Your report is available here: ", u)
		}
	}
}

func getDockerInfo() (types.Version, types.Info, error) {
	client, err := environment.Docker()
	if err != nil {
		return types.Version{}, types.Info{}, err
	}
	dockerVersion, err := client.ServerVersion(context.Background())
	if err != nil {
		return types.Version{}, types.Info{}, err
	}
	dockerInfo, err := client.Info(context.Background())
	if err != nil {
		return types.Version{}, types.Info{}, err
	}
	return dockerVersion, dockerInfo, nil
}

func uploadToHastebin(hbUrl, content string) (string, error) {
	r := strings.NewReader(content)
	u, err := url.Parse(hbUrl)
	if err != nil {
		return "", err
	}
	u.Path = path.Join(u.Path, "documents")
	res, err := http.Post(u.String(), "plain/text", r)
	if err != nil || res.StatusCode != 200 {
		fmt.Println("Failed to upload report to ", u.String(), err)
		return "", err
	}
	pres := make(map[string]interface{})
	body, err := io.ReadAll(res.Body)
	if err != nil {
		fmt.Println("Failed to parse response.", err)
		return "", err
	}
	json.Unmarshal(body, &pres)
	if key, ok := pres["key"].(string); ok {
		u, _ := url.Parse(hbUrl)
		u.Path = path.Join(u.Path, key)
		return u.String(), nil
	}
	return "", errors.New("failed to find key in response")
}

func redact(s string) string {
	if !diagnosticsArgs.IncludeEndpoints {
		return "{redacted}"
	}
	return s
}

func printHeader(w io.Writer, title string) {
	fmt.Fprintln(w, "\n|\n|", title)
	fmt.Fprintln(w, "| ------------------------------")
}
<|endoftext|>