<commit_before>// Copyright 2019 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package module

import (
	"context"
	"fmt"
	"sync"
)

// ServiceFunc 服务实际需要执行的函数
//
// 实现者需要正确处理 ctx.Done 事件，调用者可能会主动取消函数执行；
// 如果是通 ctx 取消的，应该返回其错误信息。
type ServiceFunc func(ctx context.Context) error

// ServiceState 服务的状态值
type ServiceState int8

// 几种可能的状态值
const (
	ServiceStop    ServiceState = iota + 1 // 当前处理停止状态，默认状态
	ServiceRunning                         // 正在运行
	ServiceFailed                          // 出错，不再执行后续操作
)

// Service 服务模型
type Service struct {
	ID    int // 唯一标志，在运行之后才会赋值
	Title string

	state      ServiceState
	f          ServiceFunc
	cancelFunc context.CancelFunc
	locker     sync.Mutex

	err error // 保存上次的出错内容
}

// AddService 添加新的服务
func (m *Module) AddService(f ServiceFunc, title string) {
	if m.services == nil {
		m.services = make([]*Service, 0, 5)
	}

	m.services = append(m.services, &Service{
		Title: title,
		state: ServiceStop,
		f:     f,
	})
}

// State 获取当前服务的状态
func (srv *Service) State() ServiceState {
	return srv.state
}

// Err 上次的错误信息，不会清空。
func (srv *Service) Err() error {
	return srv.err
}

// Run 开始执行该服务
func (srv *Service) Run() {
	srv.locker.Lock()
	defer srv.locker.Unlock()

	if srv.state != ServiceRunning {
		go srv.serve()
	}
}

func (srv *Service) serve() {
	defer func() {
		if msg := recover(); msg != nil {
			srv.err = fmt.Errorf("panic:%v", msg)
			srv.state = ServiceFailed
		}
	}()

	ctx := context.Background()
	ctx, srv.cancelFunc = context.WithCancel(ctx)
	srv.state = ServiceRunning

	err := srv.f(ctx)
	if err != nil && err != context.Canceled {
		srv.err = err
		srv.state = ServiceFailed
		return
	}

	srv.state = ServiceStop
}

// Stop 停止服务。
func (srv *Service) Stop() {
	srv.locker.Lock()
	defer srv.locker.Unlock()

	if srv.cancelFunc != nil {
		srv.cancelFunc()
		srv.cancelFunc = nil
	}

	srv.state = ServiceStop
}
<commit_msg>[module] 停止服务之前判断其状态<commit_after>// Copyright 2019 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package module

import (
	"context"
	"fmt"
	"sync"
)

// ServiceFunc 服务实际需要执行的函数
//
// 实现者需要正确处理 ctx.Done 事件，调用者可能会主动取消函数执行；
// 如果是通 ctx 取消的，应该返回其错误信息。
type ServiceFunc func(ctx context.Context) error

// ServiceState 服务的状态值
type ServiceState int8

// 几种可能的状态值
const (
	ServiceStop    ServiceState = iota + 1 // 当前处理停止状态，默认状态
	ServiceRunning                         // 正在运行
	ServiceFailed                          // 出错，不再执行后续操作
)

// Service 服务模型
type Service struct {
	ID    int // 唯一标志，在运行之后才会赋值
	Title string

	state      ServiceState
	f          ServiceFunc
	cancelFunc context.CancelFunc
	locker     sync.Mutex

	err error // 保存上次的出错内容
}

// AddService 添加新的服务
func (m *Module) AddService(f ServiceFunc, title string) {
	if m.services == nil {
		m.services = make([]*Service, 0, 5)
	}

	m.services = append(m.services, &Service{
		Title: title,
		state: ServiceStop,
		f:     f,
	})
}

// State 获取当前服务的状态
func (srv *Service) State() ServiceState {
	return srv.state
}

// Err 上次的错误信息，不会清空。
func (srv *Service) Err() error {
	return srv.err
}

// Run 开始执行该服务
func (srv *Service) Run() {
	srv.locker.Lock()
	defer srv.locker.Unlock()

	if srv.state != ServiceRunning {
		go srv.serve()
	}
}

func (srv *Service) serve() {
	defer func() {
		if msg := recover(); msg != nil {
			srv.err = fmt.Errorf("panic:%v", msg)
			srv.state = ServiceFailed
		}
	}()

	ctx := context.Background()
	ctx, srv.cancelFunc = context.WithCancel(ctx)
	srv.state = ServiceRunning

	err := srv.f(ctx)
	if err != nil && err != context.Canceled {
		srv.err = err
		srv.state = ServiceFailed
		return
	}

	srv.state = ServiceStop
}

// Stop 停止服务。
func (srv *Service) Stop() {
	srv.locker.Lock()
	defer srv.locker.Unlock()

	if srv.state != ServiceRunning {
		return
	}

	if srv.cancelFunc != nil {
		srv.cancelFunc()
		srv.cancelFunc = nil
	}

	srv.state = ServiceStop
}
<|endoftext|>


<commit_before>package pinger_test

import (
	"database/sql"
	"reflect"
	"strings"
	"testing"
	"time"

	_ "github.com/erikstmartin/go-testdb"
	"github.com/turnkey-commerce/go-ping-sites/database"
	"github.com/turnkey-commerce/go-ping-sites/notifier"
	"github.com/turnkey-commerce/go-ping-sites/pinger"
)

// TestNewPinger tests building the pinger object.
func TestNewPinger(t *testing.T) {
	db, _ := sql.Open("testdb", "")
	pinger.CreatePingerLog("", true)
	p := pinger.NewPinger(db, pinger.GetSitesMock, pinger.RequestURLMock,
		notifier.SendEmailMock, notifier.SendSmsMock)

	if len(p.Sites) != 3 {
		t.Fatal("Incorrect number of sites returned in new pinger.")
	}

	results, err := pinger.GetLogContent()
	if err != nil {
		t.Fatal("Failed to get log results.", err)
	}

	if !strings.Contains(results, "SITE: Test, http://www.google.com") {
		t.Fatal("Failed to load first site.")
	}
	if !strings.Contains(results, "SITE: Test 2, http://www.github.com") {
		t.Fatal("Failed to load second site.")
	}
	if !strings.Contains(results, "SITE: Test 3, http://www.test.com") {
		t.Fatal("Failed to load third site.")
	}
}

// TestStartEmptySitesPinger verifies that proper reporting is done for the case of no active sites.
func TestStartEmptySitesPinger(t *testing.T) {
	db, _ := sql.Open("testdb", "")
	pinger.CreatePingerLog("", true)
	p := pinger.NewPinger(db, pinger.GetEmptySitesMock, pinger.RequestURLMock,
		notifier.SendEmailMock, notifier.SendSmsMock)
	p.Start()

	results, err := pinger.GetLogContent()
	if err != nil {
		t.Fatal("Failed to get log results.", err)
	}

	if !strings.Contains(results, "No active sites set up for pinging in the database!") {
		t.Fatal("Failed to report empty sites.")
	}
}

// TestStartPingerErrorWithGetSites verifies that an error is handled when the get sites returns
// an error.
func TestStartPingerErrorWithGetSites(t *testing.T) {
	db, _ := sql.Open("testdb", "")
	pinger.CreatePingerLog("", true)
	p := pinger.NewPinger(db, pinger.GetSitesErrorMock, pinger.RequestURLMock,
		notifier.SendEmailMock, notifier.SendSmsMock)
	p.Start()

	results, err := pinger.GetLogContent()
	if err != nil {
		t.Fatal("Failed to get log results.", err)
	}

	if !strings.Contains(results, "Timeout accessing the SQL database.") {
		t.Fatal("Failed to report error getting the sites from the DB.")
	}
}

// TestStartAndRestartPinger starts up the pinger and then stops it after 3 seconds
func TestStartAndRestartPinger(t *testing.T) {
	// Fake db for testing.
	db, _ := sql.Open("testdb", "")
	pinger.CreatePingerLog("", true)
	p := pinger.NewPinger(db, pinger.GetSitesMock, pinger.RequestURLMock,
		notifier.SendEmailMock, notifier.SendSmsMock)
	p.Start()
	// Sleep to allow running the tests before stopping.
	time.Sleep(5 * time.Second)
	p.Stop() // Test Restart after stop
	p.Start()
	p.Stop()

	results, err := pinger.GetLogContent()
	if err != nil {
		t.Fatal("Failed to get log results.", err)
	}

	if !strings.Contains(results, "Client.Timeout") {
		t.Errorf("Failed to report timeout error.")
	}
	if !strings.Contains(results, "Test 3 Paused") {
		t.Errorf("Failed to report paused site.")
	}
	if !strings.Contains(results, "Error - HTTP Status Code") {
		t.Fatal("Failed to report bad HTTP Status Code.")
	}
	if !strings.Contains(results, "Sending Notification of Site Contacts about Test 2: Site is Down...") {
		t.Fatal("Failed to report site being down.")
	}
	if !strings.Contains(results, "Will notify status change for Test 2: Test 2 at http://www.github.com: Site is now up, response time was 300ms.") {
		t.Fatal("Failed to report change in notification.")
	}
}

// TestUpdateSiteSettings starts up the pinger and then updates the site settings.
func TestUpdateSiteSettings(t *testing.T) {
	// Fake db for testing.
	db, _ := sql.Open("testdb", "")
	pinger.CreatePingerLog("", true)
	p := pinger.NewPinger(db, pinger.GetSitesMock, pinger.RequestURLMock,
		notifier.SendEmailMock, notifier.SendSmsMock)
	p.Start()
	// Test UpdateSiteSettings
	p.UpdateSiteSettings()
	p.Stop()

	results, err := pinger.GetLogContent()
	if err != nil {
		t.Fatal("Failed to get log results.", err)
	}
	if !strings.Contains(results, "Updating the site settings due to change...") {
		t.Fatal("Failed to launch update site settings.")
	}
}

// TestUpdateSiteSettingsError tests when UpdateSiteSettings returns a DB error.
func TestUpdateSiteSettingsError(t *testing.T) {
	// Fake db for testing.
	db, _ := sql.Open("testdb", "")
	pinger.CreatePingerLog("", true)
	p := pinger.NewPinger(db, pinger.GetSitesErrorMock, pinger.RequestURLMock,
		notifier.SendEmailMock, notifier.SendSmsMock)
	p.Start()
	// Test UpdateSiteSettings with error due to database.
	err := p.UpdateSiteSettings()
	if err == nil {
		t.Error("Failed to report error with getting site settings.")
	}
	p.Stop()

	results, err := pinger.GetLogContent()
	if err != nil {
		t.Fatal("Failed to get log results.", err)
	}

	if !strings.Contains(results, "Timeout accessing the SQL database.") {
		t.Fatal("Failed to report problem updating the sites.")
	}
}

// TestInternetAccessError starts up the pinger and then stops it after 3 seconds
func TestInternetAccessError(t *testing.T) {
	// Fake db for testing.
	db, _ := sql.Open("testdb", "")
	pinger.CreatePingerLog("", true)
	p := pinger.NewPinger(db, pinger.GetSitesMock, pinger.RequestURLBadInternetAccessMock,
		notifier.SendEmailMock, notifier.SendSmsMock)
	p.Start()
	time.Sleep(1 * time.Second)
	p.Stop()

	results, err := pinger.GetLogContent()
	if err != nil {
		t.Fatal("Failed to get log results.", err)
	}

	if !strings.Contains(results, "Unable to determine site status - connect: network is unreachable") {
		t.Fatal("Failed to report Internet Access Error: ", results)
	}
}

// TestGetSites tests the database retrieval of the list of sites.
func TestGetSites(t *testing.T) {
	var sites database.Sites

	db, err := database.InitializeTestDB("")
	if err != nil {
		t.Fatal("Failed to create database:", err)
	}
	defer db.Close()

	s1 := database.Site{Name: "Test", IsActive: true, URL: "http://www.test.com",
		PingIntervalSeconds: 60, TimeoutSeconds: 30}
	err = s1.CreateSite(db)
	if err != nil {
		t.Fatal("Failed to create new site:", err)
	}

	// Create the second site.
	s2 := database.Site{Name: "Test 2", IsActive: true, URL: "http://www.example.com",
		PingIntervalSeconds: 60, TimeoutSeconds: 30}
	err = s2.CreateSite(db)
	if err != nil {
		t.Fatal("Failed to create second site:", err)
	}

	sites, err = pinger.GetSites(db)
	// Verify the first site was Loaded with proper attributes.
	if !reflect.DeepEqual(s1.URL, sites[0].URL) || !reflect.DeepEqual(s1.IsActive, sites[0].IsActive) ||
		!reflect.DeepEqual(s1.Name, sites[0].Name) || !reflect.DeepEqual(s1.PingIntervalSeconds,
		sites[0].PingIntervalSeconds) || !reflect.DeepEqual(s1.TimeoutSeconds,
		sites[0].TimeoutSeconds) || !reflect.DeepEqual(s1.SiteID, sites[0].SiteID) {
		t.Fatal("First saved site not equal to input:\n", sites[0], s1)
	}

	// Verify the second site was Loaded with proper attributes.
	if !reflect.DeepEqual(s2.URL, sites[1].URL) || !reflect.DeepEqual(s2.IsActive, sites[1].IsActive) ||
		!reflect.DeepEqual(s2.Name, sites[1].Name) || !reflect.DeepEqual(s2.PingIntervalSeconds,
		sites[1].PingIntervalSeconds) || !reflect.DeepEqual(s2.TimeoutSeconds,
		sites[1].TimeoutSeconds) || !reflect.DeepEqual(s2.SiteID, sites[1].SiteID) {
		t.Fatal("First saved site not equal to input:\n", sites[1], s2)
	}
}

func TestCreatePingerLogError(t *testing.T) {
	var logFile = "/bogusFilePath/pinger.log"
	err := pinger.CreatePingerLog(logFile, true)
	if err == nil {
		t.Error("Creation of pinger log should throw error for bad path.")
	}
}
<commit_msg>Fix timing issue with test after change in ping frequency.<commit_after>package pinger_test

import (
	"database/sql"
	"reflect"
	"strings"
	"testing"
	"time"

	_ "github.com/erikstmartin/go-testdb"
	"github.com/turnkey-commerce/go-ping-sites/database"
	"github.com/turnkey-commerce/go-ping-sites/notifier"
	"github.com/turnkey-commerce/go-ping-sites/pinger"
)

// TestNewPinger tests building the pinger object.
func TestNewPinger(t *testing.T) {
	db, _ := sql.Open("testdb", "")
	pinger.CreatePingerLog("", true)
	p := pinger.NewPinger(db, pinger.GetSitesMock, pinger.RequestURLMock,
		notifier.SendEmailMock, notifier.SendSmsMock)

	if len(p.Sites) != 3 {
		t.Fatal("Incorrect number of sites returned in new pinger.")
	}

	results, err := pinger.GetLogContent()
	if err != nil {
		t.Fatal("Failed to get log results.", err)
	}

	if !strings.Contains(results, "SITE: Test, http://www.google.com") {
		t.Fatal("Failed to load first site.")
	}
	if !strings.Contains(results, "SITE: Test 2, http://www.github.com") {
		t.Fatal("Failed to load second site.")
	}
	if !strings.Contains(results, "SITE: Test 3, http://www.test.com") {
		t.Fatal("Failed to load third site.")
	}
}

// TestStartEmptySitesPinger verifies that proper reporting is done for the case of no active sites.
func TestStartEmptySitesPinger(t *testing.T) {
	db, _ := sql.Open("testdb", "")
	pinger.CreatePingerLog("", true)
	p := pinger.NewPinger(db, pinger.GetEmptySitesMock, pinger.RequestURLMock,
		notifier.SendEmailMock, notifier.SendSmsMock)
	p.Start()

	results, err := pinger.GetLogContent()
	if err != nil {
		t.Fatal("Failed to get log results.", err)
	}

	if !strings.Contains(results, "No active sites set up for pinging in the database!") {
		t.Fatal("Failed to report empty sites.")
	}
}

// TestStartPingerErrorWithGetSites verifies that an error is handled when the get sites returns
// an error.
func TestStartPingerErrorWithGetSites(t *testing.T) {
	db, _ := sql.Open("testdb", "")
	pinger.CreatePingerLog("", true)
	p := pinger.NewPinger(db, pinger.GetSitesErrorMock, pinger.RequestURLMock,
		notifier.SendEmailMock, notifier.SendSmsMock)
	p.Start()

	results, err := pinger.GetLogContent()
	if err != nil {
		t.Fatal("Failed to get log results.", err)
	}

	if !strings.Contains(results, "Timeout accessing the SQL database.") {
		t.Fatal("Failed to report error getting the sites from the DB.")
	}
}

// TestStartAndRestartPinger starts up the pinger and then stops it after 3 seconds
func TestStartAndRestartPinger(t *testing.T) {
	// Fake db for testing.
	db, _ := sql.Open("testdb", "")
	pinger.CreatePingerLog("", true)
	p := pinger.NewPinger(db, pinger.GetSitesMock, pinger.RequestURLMock,
		notifier.SendEmailMock, notifier.SendSmsMock)
	p.Start()
	// Sleep to allow running the tests before stopping.
	time.Sleep(5 * time.Second)
	p.Stop() // Test Restart after stop
	p.Start()
	p.Stop()

	results, err := pinger.GetLogContent()
	if err != nil {
		t.Fatal("Failed to get log results.", err)
	}

	if !strings.Contains(results, "Client.Timeout") {
		t.Errorf("Failed to report timeout error.")
	}
	if !strings.Contains(results, "Test 3 Paused") {
		t.Errorf("Failed to report paused site.")
	}
	if !strings.Contains(results, "Error - HTTP Status Code") {
		t.Fatal("Failed to report bad HTTP Status Code.")
	}
	if !strings.Contains(results, "Sending Notification of Site Contacts about Test 2: Site is Down...") {
		t.Fatal("Failed to report site being down.")
	}
	if !strings.Contains(results, "Will notify status change for Test 2: Test 2 at http://www.github.com: Site is now up, response time was 300ms.") {
		t.Fatal("Failed to report change in notification.")
	}
}

// TestUpdateSiteSettings starts up the pinger and then updates the site settings.
func TestUpdateSiteSettings(t *testing.T) {
	// Fake db for testing.
	db, _ := sql.Open("testdb", "")
	pinger.CreatePingerLog("", true)
	p := pinger.NewPinger(db, pinger.GetSitesMock, pinger.RequestURLMock,
		notifier.SendEmailMock, notifier.SendSmsMock)
	p.Start()
	// Test UpdateSiteSettings
	p.UpdateSiteSettings()
	p.Stop()

	results, err := pinger.GetLogContent()
	if err != nil {
		t.Fatal("Failed to get log results.", err)
	}
	if !strings.Contains(results, "Updating the site settings due to change...") {
		t.Fatal("Failed to launch update site settings.")
	}
}

// TestUpdateSiteSettingsError tests when UpdateSiteSettings returns a DB error.
func TestUpdateSiteSettingsError(t *testing.T) {
	// Fake db for testing.
	db, _ := sql.Open("testdb", "")
	pinger.CreatePingerLog("", true)
	p := pinger.NewPinger(db, pinger.GetSitesErrorMock, pinger.RequestURLMock,
		notifier.SendEmailMock, notifier.SendSmsMock)
	p.Start()
	// Test UpdateSiteSettings with error due to database.
	err := p.UpdateSiteSettings()
	if err == nil {
		t.Error("Failed to report error with getting site settings.")
	}
	p.Stop()

	results, err := pinger.GetLogContent()
	if err != nil {
		t.Fatal("Failed to get log results.", err)
	}

	if !strings.Contains(results, "Timeout accessing the SQL database.") {
		t.Fatal("Failed to report problem updating the sites.")
	}
}

// TestInternetAccessError starts up the pinger and then stops it after 3 seconds
func TestInternetAccessError(t *testing.T) {
	// Fake db for testing.
	db, _ := sql.Open("testdb", "")
	pinger.CreatePingerLog("", true)
	p := pinger.NewPinger(db, pinger.GetSitesMock, pinger.RequestURLBadInternetAccessMock,
		notifier.SendEmailMock, notifier.SendSmsMock)
	p.Start()
	time.Sleep(5 * time.Second)
	p.Stop()

	results, err := pinger.GetLogContent()
	if err != nil {
		t.Fatal("Failed to get log results.", err)
	}

	if !strings.Contains(results, "Unable to determine site status - connect: network is unreachable") {
		t.Fatal("Failed to report Internet Access Error: ", results)
	}
}

// TestGetSites tests the database retrieval of the list of sites.
func TestGetSites(t *testing.T) {
	var sites database.Sites

	db, err := database.InitializeTestDB("")
	if err != nil {
		t.Fatal("Failed to create database:", err)
	}
	defer db.Close()

	s1 := database.Site{Name: "Test", IsActive: true, URL: "http://www.test.com",
		PingIntervalSeconds: 60, TimeoutSeconds: 30}
	err = s1.CreateSite(db)
	if err != nil {
		t.Fatal("Failed to create new site:", err)
	}

	// Create the second site.
	s2 := database.Site{Name: "Test 2", IsActive: true, URL: "http://www.example.com",
		PingIntervalSeconds: 60, TimeoutSeconds: 30}
	err = s2.CreateSite(db)
	if err != nil {
		t.Fatal("Failed to create second site:", err)
	}

	sites, err = pinger.GetSites(db)
	// Verify the first site was Loaded with proper attributes.
	if !reflect.DeepEqual(s1.URL, sites[0].URL) || !reflect.DeepEqual(s1.IsActive, sites[0].IsActive) ||
		!reflect.DeepEqual(s1.Name, sites[0].Name) || !reflect.DeepEqual(s1.PingIntervalSeconds,
		sites[0].PingIntervalSeconds) || !reflect.DeepEqual(s1.TimeoutSeconds,
		sites[0].TimeoutSeconds) || !reflect.DeepEqual(s1.SiteID, sites[0].SiteID) {
		t.Fatal("First saved site not equal to input:\n", sites[0], s1)
	}

	// Verify the second site was Loaded with proper attributes.
	if !reflect.DeepEqual(s2.URL, sites[1].URL) || !reflect.DeepEqual(s2.IsActive, sites[1].IsActive) ||
		!reflect.DeepEqual(s2.Name, sites[1].Name) || !reflect.DeepEqual(s2.PingIntervalSeconds,
		sites[1].PingIntervalSeconds) || !reflect.DeepEqual(s2.TimeoutSeconds,
		sites[1].TimeoutSeconds) || !reflect.DeepEqual(s2.SiteID, sites[1].SiteID) {
		t.Fatal("First saved site not equal to input:\n", sites[1], s2)
	}
}

func TestCreatePingerLogError(t *testing.T) {
	var logFile = "/bogusFilePath/pinger.log"
	err := pinger.CreatePingerLog(logFile, true)
	if err == nil {
		t.Error("Creation of pinger log should throw error for bad path.")
	}
}
<|endoftext|>


<commit_before>package docker

import (
	"flag"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"syscall"
	"testing"

	"github.com/goreleaser/goreleaser/config"
	"github.com/goreleaser/goreleaser/context"
	"github.com/goreleaser/goreleaser/internal/artifact"
	"github.com/goreleaser/goreleaser/pipeline"
	"github.com/stretchr/testify/assert"
)

var it = flag.Bool("it", false, "push images to docker hub")
var registry = "localhost:5000/"

func TestMain(m *testing.M) {
	if !*it {
		registry = "docker.io/"
	}
	os.Exit(m.Run())
}

func start(t *testing.T) {
	if !*it {
		return
	}
	if err := exec.Command(
		"docker", "run", "-d", "-p", "5000:5000", "--name", "registry", "registry:2",
	).Run(); err != nil {
		t.Log("failed to start docker registry", err)
		t.FailNow()
	}
}

func killAndRm(t *testing.T) {
	if !*it {
		return
	}
	t.Log("killing registry")
	_ = exec.Command("docker", "kill", "registry").Run()
	_ = exec.Command("docker", "rm", "registry").Run()
}

func TestRunPipe(t *testing.T) {
	var table = map[string]struct {
		docker  config.Docker
		publish bool
		expect  []string
		err     string
	}{
		"valid": {
			publish: true,
			docker: config.Docker{
				Image:       registry + "goreleaser/test_run_pipe",
				Goos:        "linux",
				Goarch:      "amd64",
				Dockerfile:  "testdata/Dockerfile",
				Binary:      "mybin",
				Latest:      true,
				TagTemplate: "{{.Tag}}-{{.Env.FOO}}",
				Files: []string{
					"testdata/extra_file.txt",
				},
			},
			expect: []string{
				registry + "goreleaser/test_run_pipe:v1.0.0-123",
				registry + "goreleaser/test_run_pipe:latest",
			},
			err: "",
		},
		"valid_no_latest": {
			publish: true,
			docker: config.Docker{
				Image:       registry + "goreleaser/test_run_pipe",
				Goos:        "linux",
				Goarch:      "amd64",
				Dockerfile:  "testdata/Dockerfile",
				Binary:      "mybin",
				Latest:      false,
				TagTemplate: "{{.Version}}",
				Files: []string{
					"testdata/extra_file.txt",
				},
			},
			expect: []string{
				registry + "goreleaser/test_run_pipe:1.0.0",
			},
			err: "",
		},
		"valid_dont_publish": {
			publish: false,
			docker: config.Docker{
				Image:       registry + "goreleaser/test_run_pipe",
				Goos:        "linux",
				Goarch:      "amd64",
				Dockerfile:  "testdata/Dockerfile",
				Binary:      "mybin",
				Latest:      true,
				TagTemplate: "{{.Tag}}-{{.Env.FOO}}",
				Files: []string{
					"testdata/extra_file.txt",
				},
			},
			expect: []string{
				registry + "goreleaser/test_run_pipe:v1.0.0-123",
				registry + "goreleaser/test_run_pipe:latest",
			},
			err: "",
		},
		"bad_dockerfile": {
			publish: true,
			docker: config.Docker{
				Image:       registry + "goreleaser/test_run_pipe",
				Goos:        "linux",
				Goarch:      "amd64",
				Dockerfile:  "testdata/Dockerfile.bad",
				Binary:      "mybin",
				TagTemplate: "{{.Version}}",
			},
			err: "pull access denied for nope, repository does not exist",
		},
		"template_error": {
			publish: true,
			docker: config.Docker{
				Image:       registry + "goreleaser/test_run_pipe",
				Goos:        "linux",
				Goarch:      "amd64",
				Dockerfile:  "testdata/Dockerfile",
				Binary:      "mybin",
				Latest:      true,
				TagTemplate: "{{.Tag}",
			},
			err: `template: tag:1: unexpected "}" in operand`,
		},
		"no_permissions": {
			publish: true,
			docker: config.Docker{
				Image:       "docker.io/nope",
				Goos:        "linux",
				Goarch:      "amd64",
				Binary:      "mybin",
				Dockerfile:  "testdata/Dockerfile",
				TagTemplate: "{{.Tag}}",
			},
			expect: []string{
				"docker.io/nope:latest",
				"docker.io/nope:v1.0.0",
			},
			err: `requested access to the resource is denied`,
		},
		"dockerfile_doesnt_exist": {
			publish: true,
			docker: config.Docker{
				Image:       "whatever",
				Goos:        "linux",
				Goarch:      "amd64",
				Binary:      "mybin",
				Dockerfile:  "testdata/Dockerfilezzz",
				TagTemplate: "{{.Tag}}",
			},
			err: `failed to link dockerfile`,
		},
		"extra_file_doesnt_exist": {
			publish: true,
			docker: config.Docker{
				Image:  "whatever",
				Goos:   "linux",
				Goarch: "amd64",
				Binary: "mybin",
				Files: []string{
					"testdata/nope.txt",
				},
				Dockerfile:  "testdata/Dockerfile",
				TagTemplate: "{{.Tag}}",
			},
			err: `failed to link extra file 'testdata/nope.txt'`,
		},
		"no_matching_binaries": {
			publish: true,
			docker: config.Docker{
				Image:      "whatever",
				Goos:       "darwin",
				Goarch:     "amd64",
				Binary:     "mybinnnn",
				Dockerfile: "testdata/Dockerfile",
			},
			err: "",
		},
	}

	killAndRm(t)
	start(t)
	defer killAndRm(t)

	for name, docker := range table {
		t.Run(name, func(tt *testing.T) {
			folder, err := ioutil.TempDir("", "archivetest")
			assert.NoError(t, err)
			var dist = filepath.Join(folder, "dist")
			assert.NoError(t, os.Mkdir(dist, 0755))
			assert.NoError(t, os.Mkdir(filepath.Join(dist, "mybin"), 0755))
			var binPath = filepath.Join(dist, "mybin", "mybin")
			_, err = os.Create(binPath)
			assert.NoError(t, err)

			var ctx = &context.Context{
				Version:     "1.0.0",
				Publish:     docker.publish,
				Parallelism: 4,
				Artifacts:   artifact.New(),
				Git: context.GitInfo{
					CurrentTag: "v1.0.0",
				},
				Config: config.Project{
					ProjectName: "mybin",
					Dist:        dist,
					Dockers: []config.Docker{
						docker.docker,
					},
				},
				Env: map[string]string{"FOO": "123"},
			}
			for _, os := range []string{"linux", "darwin"} {
				for _, arch := range []string{"amd64", "386"} {
					ctx.Artifacts.Add(artifact.Artifact{
						Name:   "mybin",
						Path:   binPath,
						Goarch: arch,
						Goos:   os,
						Type:   artifact.Binary,
						Extra: map[string]string{
							"Binary": "mybin",
						},
					})
				}
			}

			// this might fail as the image doesnt exist yet, so lets ignore the error
			for _, img := range docker.expect {
				_ = exec.Command("docker", "rmi", img).Run()
			}

			err = Pipe{}.Run(ctx)
			if docker.err == "" {
				assert.NoError(tt, err)
			} else {
				assert.Contains(tt, err.Error(), docker.err)
			}

			// this might should not fail as the image should have been created when
			// the step ran
			for _, img := range docker.expect {
				t.Log("removing docker image", img)
				assert.NoError(t, exec.Command("docker", "rmi", img).Run(), "could not delete image", img)
			}

		})
	}
}

func TestDescription(t *testing.T) {
	assert.NotEmpty(t, Pipe{}.String())
}

func TestNoDockers(t *testing.T) {
	assert.True(t, pipeline.IsSkip(Pipe{}.Run(context.New(config.Project{}))))
}

func TestNoDockerWithoutImageName(t *testing.T) {
	assert.True(t, pipeline.IsSkip(Pipe{}.Run(context.New(config.Project{
		Dockers: []config.Docker{
			{
				Goos: "linux",
			},
		},
	}))))
}

func TestDockerNotInPath(t *testing.T) {
	var path = os.Getenv("PATH")
	defer func() {
		assert.NoError(t, os.Setenv("PATH", path))
	}()
	assert.NoError(t, os.Setenv("PATH", ""))
	var ctx = &context.Context{
		Version: "1.0.0",
		Config: config.Project{
			Dockers: []config.Docker{
				{
					Image: "a/b",
				},
			},
		},
	}
	assert.EqualError(t, Pipe{}.Run(ctx), ErrNoDocker.Error())
}

func TestDefault(t *testing.T) {
	var ctx = &context.Context{
		Config: config.Project{
			Builds: []config.Build{
				{
					Binary: "foo",
				},
			},
			Dockers: []config.Docker{
				{
					Latest: true,
				},
			},
		},
	}
	assert.NoError(t, Pipe{}.Default(ctx))
	assert.Len(t, ctx.Config.Dockers, 1)
	var docker = ctx.Config.Dockers[0]
	assert.Equal(t, "linux", docker.Goos)
	assert.Equal(t, "amd64", docker.Goarch)
	assert.Equal(t, ctx.Config.Builds[0].Binary, docker.Binary)
	assert.Equal(t, "Dockerfile", docker.Dockerfile)
	assert.Equal(t, "{{ .Version }}", docker.TagTemplate)
}

func TestDefaultNoDockers(t *testing.T) {
	var ctx = &context.Context{
		Config: config.Project{
			Dockers: []config.Docker{},
		},
	}
	assert.NoError(t, Pipe{}.Default(ctx))
	assert.Empty(t, ctx.Config.Dockers)
}

func TestDefaultSet(t *testing.T) {
	var ctx = &context.Context{
		Config: config.Project{
			Dockers: []config.Docker{
				{
					Goos:       "windows",
					Goarch:     "i386",
					Binary:     "bar",
					Dockerfile: "Dockerfile.foo",
				},
			},
		},
	}
	assert.NoError(t, Pipe{}.Default(ctx))
	assert.Len(t, ctx.Config.Dockers, 1)
	var docker = ctx.Config.Dockers[0]
	assert.Equal(t, "windows", docker.Goos)
	assert.Equal(t, "i386", docker.Goarch)
	assert.Equal(t, "bar", docker.Binary)
	assert.Equal(t, "{{ .Version }}", docker.TagTemplate)
	assert.Equal(t, "Dockerfile.foo", docker.Dockerfile)
}

func TestLinkFile(t *testing.T) {
	const srcFile = "/tmp/test"
	const dstFile = "/tmp/linked"
	err := ioutil.WriteFile(srcFile, []byte("foo"), 0644)
	if err != nil {
		t.Log("Cannot setup test file")
		t.Fail()
	}
	err = link(srcFile, dstFile)
	if err != nil {
		t.Log("Failed to link: ", err)
		t.Fail()
	}
	if inode(srcFile) != inode(dstFile) {
		t.Log("Inodes do not match, destination file is not a link")
		t.Fail()
	}
	// cleanup
	os.Remove(srcFile)
	os.Remove(dstFile)
}

func TestLinkDirectory(t *testing.T) {
	const srcDir = "/tmp/testdir"
	const testFile = "test"
	const dstDir = "/tmp/linkedDir"

	os.Mkdir(srcDir, 0755)
	err := ioutil.WriteFile(srcDir+"/"+testFile, []byte("foo"), 0644)
	if err != nil {
		t.Log("Cannot setup test file")
		t.Fail()
	}
	err = link(srcDir, dstDir)
	if err != nil {
		t.Log("Failed to link: ", err)
		t.Fail()
	}
	if inode(srcDir+"/"+testFile) != inode(dstDir+"/"+testFile) {
		t.Log("Inodes do not match, destination file is not a link")
		t.Fail()
	}

	// cleanup
	os.RemoveAll(srcDir)
	os.RemoveAll(dstDir)
}

func TestLinkTwoLevelDirectory(t *testing.T) {
	const srcDir = "/tmp/testdir"
	const srcLevel2 = srcDir + "/level2"
	const testFile = "test"
	const dstDir = "/tmp/linkedDir"

	os.Mkdir(srcDir, 0755)
	os.Mkdir(srcLevel2, 0755)
	err := ioutil.WriteFile(srcDir+"/"+testFile, []byte("foo"), 0644)
	if err != nil {
		t.Log("Cannot setup test file")
		t.Fail()
	}
	err = ioutil.WriteFile(srcLevel2+"/"+testFile, []byte("foo"), 0644)
	if err != nil {
		t.Log("Cannot setup test file")
		t.Fail()
	}
	err = link(srcDir, dstDir)
	if err != nil {
		t.Log("Failed to link: ", err)
		t.Fail()
	}
	if inode(srcDir+"/"+testFile) != inode(dstDir+"/"+testFile) {
		t.Log("Inodes do not match")
		t.Fail()
	}
	if inode(srcLevel2+"/"+testFile) != inode(dstDir+"/level2/"+testFile) {
		t.Log("Inodes do not match")
		t.Fail()
	}
	// cleanup
	os.RemoveAll(srcDir)
	os.RemoveAll(dstDir)
}

func inode(file string) uint64 {
	fileInfo, err := os.Stat(file)
	if err != nil {
		return 0
	}
	stat := fileInfo.Sys().(*syscall.Stat_t)
	return stat.Ino
}
<commit_msg>test: flag.Parse<commit_after>package docker

import (
	"flag"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"syscall"
	"testing"

	"github.com/goreleaser/goreleaser/config"
	"github.com/goreleaser/goreleaser/context"
	"github.com/goreleaser/goreleaser/internal/artifact"
	"github.com/goreleaser/goreleaser/pipeline"
	"github.com/stretchr/testify/assert"
)

var it = flag.Bool("it", false, "push images to docker hub")
var registry = "localhost:5000/"

func TestMain(m *testing.M) {
	flag.Parse()
	if *it {
		registry = "docker.io/"
	}
	os.Exit(m.Run())
}

func start(t *testing.T) {
	if !*it {
		return
	}
	if err := exec.Command(
		"docker", "run", "-d", "-p", "5000:5000", "--name", "registry", "registry:2",
	).Run(); err != nil {
		t.Log("failed to start docker registry", err)
		t.FailNow()
	}
}

func killAndRm(t *testing.T) {
	if !*it {
		return
	}
	t.Log("killing registry")
	_ = exec.Command("docker", "kill", "registry").Run()
	_ = exec.Command("docker", "rm", "registry").Run()
}

func TestRunPipe(t *testing.T) {
	var table = map[string]struct {
		docker  config.Docker
		publish bool
		expect  []string
		err     string
	}{
		"valid": {
			publish: true,
			docker: config.Docker{
				Image:       registry + "goreleaser/test_run_pipe",
				Goos:        "linux",
				Goarch:      "amd64",
				Dockerfile:  "testdata/Dockerfile",
				Binary:      "mybin",
				Latest:      true,
				TagTemplate: "{{.Tag}}-{{.Env.FOO}}",
				Files: []string{
					"testdata/extra_file.txt",
				},
			},
			expect: []string{
				registry + "goreleaser/test_run_pipe:v1.0.0-123",
				registry + "goreleaser/test_run_pipe:latest",
			},
			err: "",
		},
		"valid_no_latest": {
			publish: true,
			docker: config.Docker{
				Image:       registry + "goreleaser/test_run_pipe",
				Goos:        "linux",
				Goarch:      "amd64",
				Dockerfile:  "testdata/Dockerfile",
				Binary:      "mybin",
				Latest:      false,
				TagTemplate: "{{.Version}}",
				Files: []string{
					"testdata/extra_file.txt",
				},
			},
			expect: []string{
				registry + "goreleaser/test_run_pipe:1.0.0",
			},
			err: "",
		},
		"valid_dont_publish": {
			publish: false,
			docker: config.Docker{
				Image:       registry + "goreleaser/test_run_pipe",
				Goos:        "linux",
				Goarch:      "amd64",
				Dockerfile:  "testdata/Dockerfile",
				Binary:      "mybin",
				Latest:      true,
				TagTemplate: "{{.Tag}}-{{.Env.FOO}}",
				Files: []string{
					"testdata/extra_file.txt",
				},
			},
			expect: []string{
				registry + "goreleaser/test_run_pipe:v1.0.0-123",
				registry + "goreleaser/test_run_pipe:latest",
			},
			err: "",
		},
		"bad_dockerfile": {
			publish: true,
			docker: config.Docker{
				Image:       registry + "goreleaser/test_run_pipe",
				Goos:        "linux",
				Goarch:      "amd64",
				Dockerfile:  "testdata/Dockerfile.bad",
				Binary:      "mybin",
				TagTemplate: "{{.Version}}",
			},
			err: "pull access denied for nope, repository does not exist",
		},
		"template_error": {
			publish: true,
			docker: config.Docker{
				Image:       registry + "goreleaser/test_run_pipe",
				Goos:        "linux",
				Goarch:      "amd64",
				Dockerfile:  "testdata/Dockerfile",
				Binary:      "mybin",
				Latest:      true,
				TagTemplate: "{{.Tag}",
			},
			err: `template: tag:1: unexpected "}" in operand`,
		},
		"no_permissions": {
			publish: true,
			docker: config.Docker{
				Image:       "docker.io/nope",
				Goos:        "linux",
				Goarch:      "amd64",
				Binary:      "mybin",
				Dockerfile:  "testdata/Dockerfile",
				TagTemplate: "{{.Tag}}",
			},
			expect: []string{
				"docker.io/nope:latest",
				"docker.io/nope:v1.0.0",
			},
			err: `requested access to the resource is denied`,
		},
		"dockerfile_doesnt_exist": {
			publish: true,
			docker: config.Docker{
				Image:       "whatever",
				Goos:        "linux",
				Goarch:      "amd64",
				Binary:      "mybin",
				Dockerfile:  "testdata/Dockerfilezzz",
				TagTemplate: "{{.Tag}}",
			},
			err: `failed to link dockerfile`,
		},
		"extra_file_doesnt_exist": {
			publish: true,
			docker: config.Docker{
				Image:  "whatever",
				Goos:   "linux",
				Goarch: "amd64",
				Binary: "mybin",
				Files: []string{
					"testdata/nope.txt",
				},
				Dockerfile:  "testdata/Dockerfile",
				TagTemplate: "{{.Tag}}",
			},
			err: `failed to link extra file 'testdata/nope.txt'`,
		},
		"no_matching_binaries": {
			publish: true,
			docker: config.Docker{
				Image:      "whatever",
				Goos:       "darwin",
				Goarch:     "amd64",
				Binary:     "mybinnnn",
				Dockerfile: "testdata/Dockerfile",
			},
			err: "",
		},
	}

	killAndRm(t)
	start(t)
	defer killAndRm(t)

	for name, docker := range table {
		t.Run(name, func(tt *testing.T) {
			folder, err := ioutil.TempDir("", "archivetest")
			assert.NoError(t, err)
			var dist = filepath.Join(folder, "dist")
			assert.NoError(t, os.Mkdir(dist, 0755))
			assert.NoError(t, os.Mkdir(filepath.Join(dist, "mybin"), 0755))
			var binPath = filepath.Join(dist, "mybin", "mybin")
			_, err = os.Create(binPath)
			assert.NoError(t, err)

			var ctx = &context.Context{
				Version:     "1.0.0",
				Publish:     docker.publish,
				Parallelism: 4,
				Artifacts:   artifact.New(),
				Git: context.GitInfo{
					CurrentTag: "v1.0.0",
				},
				Config: config.Project{
					ProjectName: "mybin",
					Dist:        dist,
					Dockers: []config.Docker{
						docker.docker,
					},
				},
				Env: map[string]string{"FOO": "123"},
			}
			for _, os := range []string{"linux", "darwin"} {
				for _, arch := range []string{"amd64", "386"} {
					ctx.Artifacts.Add(artifact.Artifact{
						Name:   "mybin",
						Path:   binPath,
						Goarch: arch,
						Goos:   os,
						Type:   artifact.Binary,
						Extra: map[string]string{
							"Binary": "mybin",
						},
					})
				}
			}

			// this might fail as the image doesnt exist yet, so lets ignore the error
			for _, img := range docker.expect {
				_ = exec.Command("docker", "rmi", img).Run()
			}

			err = Pipe{}.Run(ctx)
			if docker.err == "" {
				assert.NoError(tt, err)
			} else {
				assert.Contains(tt, err.Error(), docker.err)
			}

			// this might should not fail as the image should have been created when
			// the step ran
			for _, img := range docker.expect {
				t.Log("removing docker image", img)
				assert.NoError(t, exec.Command("docker", "rmi", img).Run(), "could not delete image", img)
			}

		})
	}
}

func TestDescription(t *testing.T) {
	assert.NotEmpty(t, Pipe{}.String())
}

func TestNoDockers(t *testing.T) {
	assert.True(t, pipeline.IsSkip(Pipe{}.Run(context.New(config.Project{}))))
}

func TestNoDockerWithoutImageName(t *testing.T) {
	assert.True(t, pipeline.IsSkip(Pipe{}.Run(context.New(config.Project{
		Dockers: []config.Docker{
			{
				Goos: "linux",
			},
		},
	}))))
}

func TestDockerNotInPath(t *testing.T) {
	var path = os.Getenv("PATH")
	defer func() {
		assert.NoError(t, os.Setenv("PATH", path))
	}()
	assert.NoError(t, os.Setenv("PATH", ""))
	var ctx = &context.Context{
		Version: "1.0.0",
		Config: config.Project{
			Dockers: []config.Docker{
				{
					Image: "a/b",
				},
			},
		},
	}
	assert.EqualError(t, Pipe{}.Run(ctx), ErrNoDocker.Error())
}

func TestDefault(t *testing.T) {
	var ctx = &context.Context{
		Config: config.Project{
			Builds: []config.Build{
				{
					Binary: "foo",
				},
			},
			Dockers: []config.Docker{
				{
					Latest: true,
				},
			},
		},
	}
	assert.NoError(t, Pipe{}.Default(ctx))
	assert.Len(t, ctx.Config.Dockers, 1)
	var docker = ctx.Config.Dockers[0]
	assert.Equal(t, "linux", docker.Goos)
	assert.Equal(t, "amd64", docker.Goarch)
	assert.Equal(t, ctx.Config.Builds[0].Binary, docker.Binary)
	assert.Equal(t, "Dockerfile", docker.Dockerfile)
	assert.Equal(t, "{{ .Version }}", docker.TagTemplate)
}

func TestDefaultNoDockers(t *testing.T) {
	var ctx = &context.Context{
		Config: config.Project{
			Dockers: []config.Docker{},
		},
	}
	assert.NoError(t, Pipe{}.Default(ctx))
	assert.Empty(t, ctx.Config.Dockers)
}

func TestDefaultSet(t *testing.T) {
	var ctx = &context.Context{
		Config: config.Project{
			Dockers: []config.Docker{
				{
					Goos:       "windows",
					Goarch:     "i386",
					Binary:     "bar",
					Dockerfile: "Dockerfile.foo",
				},
			},
		},
	}
	assert.NoError(t, Pipe{}.Default(ctx))
	assert.Len(t, ctx.Config.Dockers, 1)
	var docker = ctx.Config.Dockers[0]
	assert.Equal(t, "windows", docker.Goos)
	assert.Equal(t, "i386", docker.Goarch)
	assert.Equal(t, "bar", docker.Binary)
	assert.Equal(t, "{{ .Version }}", docker.TagTemplate)
	assert.Equal(t, "Dockerfile.foo", docker.Dockerfile)
}

func TestLinkFile(t *testing.T) {
	const srcFile = "/tmp/test"
	const dstFile = "/tmp/linked"
	err := ioutil.WriteFile(srcFile, []byte("foo"), 0644)
	if err != nil {
		t.Log("Cannot setup test file")
		t.Fail()
	}
	err = link(srcFile, dstFile)
	if err != nil {
		t.Log("Failed to link: ", err)
		t.Fail()
	}
	if inode(srcFile) != inode(dstFile) {
		t.Log("Inodes do not match, destination file is not a link")
		t.Fail()
	}
	// cleanup
	os.Remove(srcFile)
	os.Remove(dstFile)
}

func TestLinkDirectory(t *testing.T) {
	const srcDir = "/tmp/testdir"
	const testFile = "test"
	const dstDir = "/tmp/linkedDir"

	os.Mkdir(srcDir, 0755)
	err := ioutil.WriteFile(srcDir+"/"+testFile, []byte("foo"), 0644)
	if err != nil {
		t.Log("Cannot setup test file")
		t.Fail()
	}
	err = link(srcDir, dstDir)
	if err != nil {
		t.Log("Failed to link: ", err)
		t.Fail()
	}
	if inode(srcDir+"/"+testFile) != inode(dstDir+"/"+testFile) {
		t.Log("Inodes do not match, destination file is not a link")
		t.Fail()
	}

	// cleanup
	os.RemoveAll(srcDir)
	os.RemoveAll(dstDir)
}

func TestLinkTwoLevelDirectory(t *testing.T) {
	const srcDir = "/tmp/testdir"
	const srcLevel2 = srcDir + "/level2"
	const testFile = "test"
	const dstDir = "/tmp/linkedDir"

	os.Mkdir(srcDir, 0755)
	os.Mkdir(srcLevel2, 0755)
	err := ioutil.WriteFile(srcDir+"/"+testFile, []byte("foo"), 0644)
	if err != nil {
		t.Log("Cannot setup test file")
		t.Fail()
	}
	err = ioutil.WriteFile(srcLevel2+"/"+testFile, []byte("foo"), 0644)
	if err != nil {
		t.Log("Cannot setup test file")
		t.Fail()
	}
	err = link(srcDir, dstDir)
	if err != nil {
		t.Log("Failed to link: ", err)
		t.Fail()
	}
	if inode(srcDir+"/"+testFile) != inode(dstDir+"/"+testFile) {
		t.Log("Inodes do not match")
		t.Fail()
	}
	if inode(srcLevel2+"/"+testFile) != inode(dstDir+"/level2/"+testFile) {
		t.Log("Inodes do not match")
		t.Fail()
	}
	// cleanup
	os.RemoveAll(srcDir)
	os.RemoveAll(dstDir)
}

func inode(file string) uint64 {
	fileInfo, err := os.Stat(file)
	if err != nil {
		return 0
	}
	stat := fileInfo.Sys().(*syscall.Stat_t)
	return stat.Ino
}
<|endoftext|>


<commit_before>package compression

import (
	"bytes"
	"compress/bzip2"
	"fmt"
	"io"
	"io/ioutil"

	"github.com/containers/image/v5/pkg/compression/internal"
	"github.com/containers/image/v5/pkg/compression/types"
	"github.com/containers/storage/pkg/chunked"
	"github.com/klauspost/pgzip"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/ulikunitz/xz"
)

// Algorithm is a compression algorithm that can be used for CompressStream.
type Algorithm = types.Algorithm

var (
	// Gzip compression.
	Gzip = internal.NewAlgorithm("gzip", "gzip",
		[]byte{0x1F, 0x8B, 0x08}, GzipDecompressor, gzipCompressor)
	// Bzip2 compression.
	Bzip2 = internal.NewAlgorithm("bzip2", "bzip2",
		[]byte{0x42, 0x5A, 0x68}, Bzip2Decompressor, bzip2Compressor)
	// Xz compression.
	Xz = internal.NewAlgorithm("Xz", "xz",
		[]byte{0xFD, 0x37, 0x7A, 0x58, 0x5A, 0x00}, XzDecompressor, xzCompressor)
	// Zstd compression.
	Zstd = internal.NewAlgorithm("zstd", "zstd",
		[]byte{0x28, 0xb5, 0x2f, 0xfd}, ZstdDecompressor, zstdCompressor)
	// Zstd:chunked compression.
	ZstdChunked = internal.NewAlgorithm("zstd:chunked", "zstd",
		nil, ZstdDecompressor, chunked.ZstdCompressor)

	compressionAlgorithms = map[string]Algorithm{
		Gzip.Name():        Gzip,
		Bzip2.Name():       Bzip2,
		Xz.Name():          Xz,
		Zstd.Name():        Zstd,
		ZstdChunked.Name(): ZstdChunked,
	}
)

// AlgorithmByName returns the compressor by its name
func AlgorithmByName(name string) (Algorithm, error) {
	algorithm, ok := compressionAlgorithms[name]
	if ok {
		return algorithm, nil
	}
	return Algorithm{}, fmt.Errorf("cannot find compressor for %q", name)
}

// DecompressorFunc returns the decompressed stream, given a compressed stream.
// The caller must call Close() on the decompressed stream (even if the compressed input stream does not need closing!).
type DecompressorFunc = internal.DecompressorFunc

// GzipDecompressor is a DecompressorFunc for the gzip compression algorithm.
func GzipDecompressor(r io.Reader) (io.ReadCloser, error) {
	return pgzip.NewReader(r)
}

// Bzip2Decompressor is a DecompressorFunc for the bzip2 compression algorithm.
func Bzip2Decompressor(r io.Reader) (io.ReadCloser, error) {
	return ioutil.NopCloser(bzip2.NewReader(r)), nil
}

// XzDecompressor is a DecompressorFunc for the xz compression algorithm.
func XzDecompressor(r io.Reader) (io.ReadCloser, error) {
	r, err := xz.NewReader(r)
	if err != nil {
		return nil, err
	}
	return ioutil.NopCloser(r), nil
}

// gzipCompressor is a CompressorFunc for the gzip compression algorithm.
func gzipCompressor(r io.Writer, metadata map[string]string, level *int) (io.WriteCloser, error) {
	if level != nil {
		return pgzip.NewWriterLevel(r, *level)
	}
	return pgzip.NewWriter(r), nil
}

// bzip2Compressor is a CompressorFunc for the bzip2 compression algorithm.
func bzip2Compressor(r io.Writer, metadata map[string]string, level *int) (io.WriteCloser, error) {
	return nil, fmt.Errorf("bzip2 compression not supported")
}

// xzCompressor is a CompressorFunc for the xz compression algorithm.
func xzCompressor(r io.Writer, metadata map[string]string, level *int) (io.WriteCloser, error) {
	return xz.NewWriter(r)
}

// CompressStream returns the compressor by its name
func CompressStream(dest io.Writer, algo Algorithm, level *int) (io.WriteCloser, error) {
	m := map[string]string{}
	return internal.AlgorithmCompressor(algo)(dest, m, level)
}

// CompressStreamWithMetadata returns the compressor by its name.  If the compression
// generates any metadata, it is written to the provided metadata map.
func CompressStreamWithMetadata(dest io.Writer, metadata map[string]string, algo Algorithm, level *int) (io.WriteCloser, error) {
	return internal.AlgorithmCompressor(algo)(dest, metadata, level)
}

// DetectCompressionFormat returns an Algorithm and DecompressorFunc if the input is recognized as a compressed format, an invalid
// value and nil otherwise.
// Because it consumes the start of input, other consumers must use the returned io.Reader instead to also read from the beginning.
func DetectCompressionFormat(input io.Reader) (Algorithm, DecompressorFunc, io.Reader, error) {
	buffer := [8]byte{}

	n, err := io.ReadAtLeast(input, buffer[:], len(buffer))
	if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
		// This is a “real” error. We could just ignore it this time, process the data we have, and hope that the source will report the same error again.
		// Instead, fail immediately with the original error cause instead of a possibly secondary/misleading error returned later.
		return Algorithm{}, nil, nil, err
	}

	var retAlgo Algorithm
	var decompressor DecompressorFunc
	for _, algo := range compressionAlgorithms {
		prefix := internal.AlgorithmPrefix(algo)
		if len(prefix) > 0 && bytes.HasPrefix(buffer[:n], prefix) {
			logrus.Debugf("Detected compression format %s", algo.Name())
			retAlgo = algo
			decompressor = internal.AlgorithmDecompressor(algo)
			break
		}
	}
	if decompressor == nil {
		logrus.Debugf("No compression detected")
	}

	return retAlgo, decompressor, io.MultiReader(bytes.NewReader(buffer[:n]), input), nil
}

// DetectCompression returns a DecompressorFunc if the input is recognized as a compressed format, nil otherwise.
// Because it consumes the start of input, other consumers must use the returned io.Reader instead to also read from the beginning.
func DetectCompression(input io.Reader) (DecompressorFunc, io.Reader, error) {
	_, d, r, e := DetectCompressionFormat(input)
	return d, r, e
}

// AutoDecompress takes a stream and returns an uncompressed version of the
// same stream.
// The caller must call Close() on the returned stream (even if the input does not need,
// or does not even support, closing!).
func AutoDecompress(stream io.Reader) (io.ReadCloser, bool, error) {
	decompressor, stream, err := DetectCompression(stream)
	if err != nil {
		return nil, false, errors.Wrapf(err, "detecting compression")
	}
	var res io.ReadCloser
	if decompressor != nil {
		res, err = decompressor(stream)
		if err != nil {
			return nil, false, errors.Wrapf(err, "initializing decompression")
		}
	} else {
		res = ioutil.NopCloser(stream)
	}
	return res, decompressor != nil, nil
}
<commit_msg>Don't change InternalUnstableUndocumentedMIMEQuestionMark of Xz<commit_after>package compression

import (
	"bytes"
	"compress/bzip2"
	"fmt"
	"io"
	"io/ioutil"

	"github.com/containers/image/v5/pkg/compression/internal"
	"github.com/containers/image/v5/pkg/compression/types"
	"github.com/containers/storage/pkg/chunked"
	"github.com/klauspost/pgzip"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
	"github.com/ulikunitz/xz"
)

// Algorithm is a compression algorithm that can be used for CompressStream.
type Algorithm = types.Algorithm

var (
	// Gzip compression.
	Gzip = internal.NewAlgorithm("gzip", "gzip",
		[]byte{0x1F, 0x8B, 0x08}, GzipDecompressor, gzipCompressor)
	// Bzip2 compression.
	Bzip2 = internal.NewAlgorithm("bzip2", "bzip2",
		[]byte{0x42, 0x5A, 0x68}, Bzip2Decompressor, bzip2Compressor)
	// Xz compression.
	Xz = internal.NewAlgorithm("Xz", "Xz",
		[]byte{0xFD, 0x37, 0x7A, 0x58, 0x5A, 0x00}, XzDecompressor, xzCompressor)
	// Zstd compression.
	Zstd = internal.NewAlgorithm("zstd", "zstd",
		[]byte{0x28, 0xb5, 0x2f, 0xfd}, ZstdDecompressor, zstdCompressor)
	// Zstd:chunked compression.
	ZstdChunked = internal.NewAlgorithm("zstd:chunked", "zstd",
		nil, ZstdDecompressor, chunked.ZstdCompressor)

	compressionAlgorithms = map[string]Algorithm{
		Gzip.Name():        Gzip,
		Bzip2.Name():       Bzip2,
		Xz.Name():          Xz,
		Zstd.Name():        Zstd,
		ZstdChunked.Name(): ZstdChunked,
	}
)

// AlgorithmByName returns the compressor by its name
func AlgorithmByName(name string) (Algorithm, error) {
	algorithm, ok := compressionAlgorithms[name]
	if ok {
		return algorithm, nil
	}
	return Algorithm{}, fmt.Errorf("cannot find compressor for %q", name)
}

// DecompressorFunc returns the decompressed stream, given a compressed stream.
// The caller must call Close() on the decompressed stream (even if the compressed input stream does not need closing!).
type DecompressorFunc = internal.DecompressorFunc

// GzipDecompressor is a DecompressorFunc for the gzip compression algorithm.
func GzipDecompressor(r io.Reader) (io.ReadCloser, error) {
	return pgzip.NewReader(r)
}

// Bzip2Decompressor is a DecompressorFunc for the bzip2 compression algorithm.
func Bzip2Decompressor(r io.Reader) (io.ReadCloser, error) {
	return ioutil.NopCloser(bzip2.NewReader(r)), nil
}

// XzDecompressor is a DecompressorFunc for the xz compression algorithm.
func XzDecompressor(r io.Reader) (io.ReadCloser, error) {
	r, err := xz.NewReader(r)
	if err != nil {
		return nil, err
	}
	return ioutil.NopCloser(r), nil
}

// gzipCompressor is a CompressorFunc for the gzip compression algorithm.
func gzipCompressor(r io.Writer, metadata map[string]string, level *int) (io.WriteCloser, error) {
	if level != nil {
		return pgzip.NewWriterLevel(r, *level)
	}
	return pgzip.NewWriter(r), nil
}

// bzip2Compressor is a CompressorFunc for the bzip2 compression algorithm.
func bzip2Compressor(r io.Writer, metadata map[string]string, level *int) (io.WriteCloser, error) {
	return nil, fmt.Errorf("bzip2 compression not supported")
}

// xzCompressor is a CompressorFunc for the xz compression algorithm.
func xzCompressor(r io.Writer, metadata map[string]string, level *int) (io.WriteCloser, error) {
	return xz.NewWriter(r)
}

// CompressStream returns the compressor by its name
func CompressStream(dest io.Writer, algo Algorithm, level *int) (io.WriteCloser, error) {
	m := map[string]string{}
	return internal.AlgorithmCompressor(algo)(dest, m, level)
}

// CompressStreamWithMetadata returns the compressor by its name.  If the compression
// generates any metadata, it is written to the provided metadata map.
func CompressStreamWithMetadata(dest io.Writer, metadata map[string]string, algo Algorithm, level *int) (io.WriteCloser, error) {
	return internal.AlgorithmCompressor(algo)(dest, metadata, level)
}

// DetectCompressionFormat returns an Algorithm and DecompressorFunc if the input is recognized as a compressed format, an invalid
// value and nil otherwise.
// Because it consumes the start of input, other consumers must use the returned io.Reader instead to also read from the beginning.
func DetectCompressionFormat(input io.Reader) (Algorithm, DecompressorFunc, io.Reader, error) {
	buffer := [8]byte{}

	n, err := io.ReadAtLeast(input, buffer[:], len(buffer))
	if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
		// This is a “real” error. We could just ignore it this time, process the data we have, and hope that the source will report the same error again.
		// Instead, fail immediately with the original error cause instead of a possibly secondary/misleading error returned later.
		return Algorithm{}, nil, nil, err
	}

	var retAlgo Algorithm
	var decompressor DecompressorFunc
	for _, algo := range compressionAlgorithms {
		prefix := internal.AlgorithmPrefix(algo)
		if len(prefix) > 0 && bytes.HasPrefix(buffer[:n], prefix) {
			logrus.Debugf("Detected compression format %s", algo.Name())
			retAlgo = algo
			decompressor = internal.AlgorithmDecompressor(algo)
			break
		}
	}
	if decompressor == nil {
		logrus.Debugf("No compression detected")
	}

	return retAlgo, decompressor, io.MultiReader(bytes.NewReader(buffer[:n]), input), nil
}

// DetectCompression returns a DecompressorFunc if the input is recognized as a compressed format, nil otherwise.
// Because it consumes the start of input, other consumers must use the returned io.Reader instead to also read from the beginning.
func DetectCompression(input io.Reader) (DecompressorFunc, io.Reader, error) {
	_, d, r, e := DetectCompressionFormat(input)
	return d, r, e
}

// AutoDecompress takes a stream and returns an uncompressed version of the
// same stream.
// The caller must call Close() on the returned stream (even if the input does not need,
// or does not even support, closing!).
func AutoDecompress(stream io.Reader) (io.ReadCloser, bool, error) {
	decompressor, stream, err := DetectCompression(stream)
	if err != nil {
		return nil, false, errors.Wrapf(err, "detecting compression")
	}
	var res io.ReadCloser
	if decompressor != nil {
		res, err = decompressor(stream)
		if err != nil {
			return nil, false, errors.Wrapf(err, "initializing decompression")
		}
	} else {
		res = ioutil.NopCloser(stream)
	}
	return res, decompressor != nil, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Kubernetes Authors All rights reserved.

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 oci

import (
	"bufio"
	"bytes"
	"fmt"
	"os/exec"
	"strings"

	"github.com/golang/glog"
	"github.com/pkg/errors"
)

// DeleteAllVolumesByLabel deletes all volumes that have a specific label
// if there is no volume to delete it will return nil
func DeleteAllVolumesByLabel(ociBin string, label string) []error {
	var deleteErrs []error
	glog.Infof("trying to delete all %s volumes with label %s", ociBin, label)

	vs, err := allVolumesByLabel(ociBin, label)

	if err != nil {
		return []error{fmt.Errorf("listing volumes by label %q: %v", label, err)}
	}

	for _, v := range vs {
		if _, err := WarnIfSlow(ociBin, "volume", "rm", "--force", v); err != nil {
			deleteErrs = append(deleteErrs, fmt.Errorf("deleting %q", v))
		}
	}

	return deleteErrs
}

// PruneAllVolumesByLabel deletes all volumes that have a specific label
// if there is no volume to delete it will return nil
// example: docker volume prune -f --filter label=name.minikube.sigs.k8s.io=minikube
func PruneAllVolumesByLabel(ociBin string, label string) []error {
	var deleteErrs []error
	glog.Infof("trying to prune all %s volumes with label %s", ociBin, label)

	if _, err := WarnIfSlow(ociBin, "volume", "prune", "-f", "--filter", "label="+label); err != nil {
		deleteErrs = append(deleteErrs, errors.Wrapf(err, "prune volume by label %s", label))
	}

	return deleteErrs
}

// allVolumesByLabel returns name of all docker volumes by a specific label
// will not return error if there is no volume found.
func allVolumesByLabel(ociBin string, label string) ([]string, error) {
	rr, err := cli.RunCmd(exec.Command(ociBin, "volume", "ls", "--filter", "label="+label, "--format", "{{.Name}}"))
	s := bufio.NewScanner(bytes.NewReader(rr.Stdout.Bytes()))
	var vols []string
	for s.Scan() {
		v := strings.TrimSpace(s.Text())
		if v != "" {
			vols = append(vols, v)
		}
	}
	return vols, err
}

// ExtractTarballToVolume runs a docker image imageName which extracts the tarball at tarballPath
// to the volume named volumeName
func ExtractTarballToVolume(tarballPath, volumeName, imageName string) error {
	cmd := exec.Command(Docker, "run", "--rm", "--entrypoint", "/usr/bin/tar", "-v", fmt.Sprintf("%s:/preloaded.tar:ro", tarballPath), "-v", fmt.Sprintf("%s:/extractDir", volumeName), imageName, "-I", "lz4", "-xvf", "/preloaded.tar", "-C", "/extractDir")
	if _, err := cli.RunCmd(cmd); err != nil {
		return err
	}
	return nil
}

// createDockerVolume creates a docker volume to be attached to the container with correct labels and prefixes based on profile name
// Caution ! if volume already exists does NOT return an error and will not apply the minikube labels on it.
// TODO: this should be fixed as a part of https://github.com/kubernetes/minikube/issues/6530
func createDockerVolume(profile string, nodeName string) error {
	if _, err := cli.RunCmd(exec.Command(Docker, "volume", "create", nodeName, "--label", fmt.Sprintf("%s=%s", ProfileLabelKey, profile), "--label", fmt.Sprintf("%s=%s", CreatedByLabelKey, "true"))); err != nil {
		return err
	}
	return nil
}
<commit_msg>unexport warnIfSlow func<commit_after>/*
Copyright 2020 The Kubernetes Authors All rights reserved.

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 oci

import (
	"bufio"
	"bytes"
	"fmt"
	"os/exec"
	"strings"

	"github.com/golang/glog"
	"github.com/pkg/errors"
)

// DeleteAllVolumesByLabel deletes all volumes that have a specific label
// if there is no volume to delete it will return nil
func DeleteAllVolumesByLabel(ociBin string, label string) []error {
	var deleteErrs []error
	glog.Infof("trying to delete all %s volumes with label %s", ociBin, label)

	vs, err := allVolumesByLabel(ociBin, label)

	if err != nil {
		return []error{fmt.Errorf("listing volumes by label %q: %v", label, err)}
	}

	for _, v := range vs {
		if _, err := warnIfSlow(ociBin, "volume", "rm", "--force", v); err != nil {
			deleteErrs = append(deleteErrs, fmt.Errorf("deleting %q", v))
		}
	}

	return deleteErrs
}

// PruneAllVolumesByLabel deletes all volumes that have a specific label
// if there is no volume to delete it will return nil
// example: docker volume prune -f --filter label=name.minikube.sigs.k8s.io=minikube
func PruneAllVolumesByLabel(ociBin string, label string) []error {
	var deleteErrs []error
	glog.Infof("trying to prune all %s volumes with label %s", ociBin, label)

	if _, err := warnIfSlow(ociBin, "volume", "prune", "-f", "--filter", "label="+label); err != nil {
		deleteErrs = append(deleteErrs, errors.Wrapf(err, "prune volume by label %s", label))
	}

	return deleteErrs
}

// allVolumesByLabel returns name of all docker volumes by a specific label
// will not return error if there is no volume found.
func allVolumesByLabel(ociBin string, label string) ([]string, error) {
	rr, err := cli.RunCmd(exec.Command(ociBin, "volume", "ls", "--filter", "label="+label, "--format", "{{.Name}}"))
	s := bufio.NewScanner(bytes.NewReader(rr.Stdout.Bytes()))
	var vols []string
	for s.Scan() {
		v := strings.TrimSpace(s.Text())
		if v != "" {
			vols = append(vols, v)
		}
	}
	return vols, err
}

// ExtractTarballToVolume runs a docker image imageName which extracts the tarball at tarballPath
// to the volume named volumeName
func ExtractTarballToVolume(tarballPath, volumeName, imageName string) error {
	cmd := exec.Command(Docker, "run", "--rm", "--entrypoint", "/usr/bin/tar", "-v", fmt.Sprintf("%s:/preloaded.tar:ro", tarballPath), "-v", fmt.Sprintf("%s:/extractDir", volumeName), imageName, "-I", "lz4", "-xvf", "/preloaded.tar", "-C", "/extractDir")
	if _, err := cli.RunCmd(cmd); err != nil {
		return err
	}
	return nil
}

// createDockerVolume creates a docker volume to be attached to the container with correct labels and prefixes based on profile name
// Caution ! if volume already exists does NOT return an error and will not apply the minikube labels on it.
// TODO: this should be fixed as a part of https://github.com/kubernetes/minikube/issues/6530
func createDockerVolume(profile string, nodeName string) error {
	if _, err := cli.RunCmd(exec.Command(Docker, "volume", "create", nodeName, "--label", fmt.Sprintf("%s=%s", ProfileLabelKey, profile), "--label", fmt.Sprintf("%s=%s", CreatedByLabelKey, "true"))); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2021 Authors of Cilium
//
// 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 recorder

import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"os"
	"path"
	"regexp"
	"time"

	recorderpb "github.com/cilium/cilium/api/v1/recorder"
	"github.com/cilium/cilium/pkg/cidr"
	"github.com/cilium/cilium/pkg/hubble/recorder/pcap"
	"github.com/cilium/cilium/pkg/hubble/recorder/recorderoption"
	"github.com/cilium/cilium/pkg/hubble/recorder/sink"
	"github.com/cilium/cilium/pkg/idpool"
	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"
	nodeTypes "github.com/cilium/cilium/pkg/node/types"
	"github.com/cilium/cilium/pkg/recorder"
	"github.com/cilium/cilium/pkg/u8proto"

	"github.com/sirupsen/logrus"
	"google.golang.org/protobuf/types/known/timestamppb"
)

var log = logging.DefaultLogger.WithField(logfields.LogSubsys, "hubble-recorder")

var _ recorderpb.RecorderServer = (*Service)(nil)

const (
	minRuleID = 1
	maxRuleID = 65534

	defaultFileSinkPrefix = "hubble"
)

type Service struct {
	recorder *recorder.Recorder
	dispatch *sink.Dispatch
	ruleIDs  idpool.IDPool
	opts     recorderoption.Options
}

type recording struct {
	ruleID   uint16
	filePath string
	handle   *sink.Handle
}

func NewService(r *recorder.Recorder, d *sink.Dispatch, options ...recorderoption.Option) (*Service, error) {
	opts := recorderoption.Default
	for _, o := range options {
		if err := o(&opts); err != nil {
			return nil, err
		}
	}

	if len(opts.StoragePath) == 0 {
		return nil, errors.New("storage path must not be empty")
	}

	if err := os.MkdirAll(opts.StoragePath, 0600); err != nil {
		return nil, fmt.Errorf("failed to create storage path directory: %w", err)
	}

	return &Service{
		recorder: r,
		dispatch: d,
		ruleIDs:  idpool.NewIDPool(minRuleID, maxRuleID),
		opts:     opts,
	}, nil
}

func (s *Service) Record(stream recorderpb.Recorder_RecordServer) error {
	ctx, cancel := context.WithCancel(stream.Context())
	defer cancel()

	req, err := stream.Recv()
	if err != nil {
		return err
	}
	startRecording := req.GetStart()
	if startRecording == nil {
		return fmt.Errorf("received invalid request %q, expected start request", req)
	}

	rec, err := s.startRecording(ctx, startRecording)
	if err != nil {
		return err
	}

	go s.watchRecording(ctx, rec.handle, stream)

	req, err = stream.Recv()
	if err != nil {
		return err
	}

	if req.GetStop() == nil {
		return fmt.Errorf("received invalid request %q, expected stop request", req)
	}

	resp, err := s.stopRecording(ctx, rec)
	if err != nil {
		return err
	}

	err = stream.Send(&recorderpb.RecordResponse{
		NodeName: nodeTypes.GetAbsoluteNodeName(),
		Time:     timestamppb.Now(),
		ResponseType: &recorderpb.RecordResponse_Stopped{
			Stopped: resp,
		},
	})
	if err != nil {
		return err
	}

	return nil
}

const fileExistsRetries = 100

func createPcapFile(basedir, prefix string) (f *os.File, filePath string, err error) {
	try := 0
	for {
		startTime := time.Now().Unix()
		random := rand.Uint32()
		nodeName := nodeTypes.GetAbsoluteNodeName()
		name := fmt.Sprintf("%s_%d_%d_%s.pcap", prefix, startTime, random, nodeName)
		filePath = path.Join(basedir, name)
		f, err = os.OpenFile(filePath, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0600)
		if err != nil {
			if os.IsExist(err) {
				if try++; try < fileExistsRetries {
					continue
				}
			}
			return f, "", fmt.Errorf("failed to create pcap file %q: %w", filePath, err)
		}

		return f, filePath, nil
	}
}

func parseFilters(include []*recorderpb.Filter) ([]recorder.RecorderTuple, error) {
	if len(include) == 0 {
		return nil, errors.New("need to specify at least one include filter")
	}

	filters := []recorder.RecorderTuple{}
	for _, f := range include {
		srcIP, srcPrefix, err := net.ParseCIDR(f.GetSourceCidr())
		if err != nil {
			return nil, fmt.Errorf("failed to parse source cidr %q: %w", f.GetSourceCidr(), err)
		}

		dstIP, dstPrefix, err := net.ParseCIDR(f.GetDestinationCidr())
		if err != nil {
			return nil, fmt.Errorf("failed to parse source cidr %q: %w", f.GetDestinationCidr(), err)
		}

		if (srcIP.To4() == nil) != (dstIP.To4() == nil) {
			return nil, fmt.Errorf("source (%s) and destination cidr (%s) must be same protocol version",
				f.GetSourceCidr(), f.GetDestinationCidr())
		}

		const maxPort = 65535
		if f.GetSourcePort() > maxPort {
			return nil, fmt.Errorf("source port %d out of range", f.GetSourcePort())
		}

		if f.GetDestinationPort() > maxPort {
			return nil, fmt.Errorf("destination port %d out of range", f.GetDestinationPort())
		}

		filters = append(filters, recorder.RecorderTuple{
			SrcPrefix: *cidr.NewCIDR(srcPrefix),
			SrcPort:   uint16(f.GetSourcePort()),
			DstPrefix: *cidr.NewCIDR(dstPrefix),
			DstPort:   uint16(f.GetDestinationPort()),
			Proto:     u8proto.U8proto(f.GetProtocol()),
		})
	}

	return filters, nil
}

var fileSinkPrefixRegex = regexp.MustCompile("^[a-z][a-z0-9]{0,19}$")

func (s *Service) startRecording(ctx context.Context, req *recorderpb.StartRecording) (*recording, error) {
	capLen := req.GetMaxCaptureLength()
	prefix := req.GetFilesink().GetFilePrefix()
	if prefix == "" {
		prefix = defaultFileSinkPrefix
	}

	if !fileSinkPrefixRegex.MatchString(prefix) {
		return nil, fmt.Errorf("invalid file sink prefix: %q", prefix)
	}

	filters, err := parseFilters(req.GetInclude())
	if err != nil {
		return nil, err
	}

	leaseID := s.ruleIDs.LeaseAvailableID()
	ruleID := uint16(leaseID)
	if leaseID == idpool.NoID {
		return nil, errors.New("unable to allocate capture rule id")
	}

	f, filePath, err := createPcapFile(s.opts.StoragePath, prefix)
	if err != nil {
		return nil, err
	}

	r := &recording{ruleID: ruleID, filePath: filePath}
	defer func() {
		// clean up the recording if any of the subsequent steps fails
		if err != nil {
			_, _ = s.recorder.DeleteRecorder(recorder.ID(r.ruleID))
			// remove the created pcap file
			_ = f.Close()
			_ = os.Remove(r.filePath)
			// release will also invalidate the lease
			_ = s.ruleIDs.Release(idpool.ID(r.ruleID))
		}
	}()

	log.WithFields(logrus.Fields{
		"ruleID":   r.ruleID,
		"filePath": r.filePath,
	}).Debug("starting new recording")

	config := sink.PcapSink{
		RuleID: ruleID,
		Header: pcap.Header{
			SnapshotLength: capLen,
			Datalink:       pcap.Ethernet,
		},
		Writer: pcap.NewWriter(f),
	}

	r.handle, err = s.dispatch.StartSink(ctx, config)
	if err != nil {
		return nil, err
	}

	recInfo := &recorder.RecInfo{
		ID:      recorder.ID(ruleID),
		CapLen:  uint16(capLen),
		Filters: filters,
	}
	_, err = s.recorder.UpsertRecorder(recInfo)
	if err != nil {
		return nil, err
	}

	s.ruleIDs.Use(leaseID)

	return r, nil
}

func (s *Service) stopRecording(ctx context.Context, r *recording) (*recorderpb.RecordingStoppedResponse, error) {
	log.WithFields(logrus.Fields{
		"ruleID":   r.ruleID,
		"filePath": r.filePath,
	}).Debug("stopping recording")

	_, err := s.recorder.DeleteRecorder(recorder.ID(r.ruleID))
	if err != nil {
		return nil, err
	}

	r.handle.Stop()
	select {
	case <-r.handle.Done:
		if err = r.handle.Err(); err != nil {
			return nil, err
		}
	case <-ctx.Done():
		return nil, fmt.Errorf("timed out waiting for sink to close: %w", ctx.Err())
	}
	stats := r.handle.Stats()

	s.ruleIDs.Release(idpool.ID(r.ruleID))

	return &recorderpb.RecordingStoppedResponse{
		Stats: &recorderpb.RecordingStatistics{
			BytesCaptured:   stats.BytesWritten,
			PacketsCaptured: stats.PacketsWritten,
			BytesLost:       stats.BytesLost,
			PacketsLost:     stats.PacketsLost,
		},
		Filesink: &recorderpb.FileSinkResult{FilePath: r.filePath},
	}, nil
}

func (s *Service) watchRecording(ctx context.Context, h *sink.Handle, stream recorderpb.Recorder_RecordServer) {
	for {
		stats := h.Stats()
		err := stream.Send(&recorderpb.RecordResponse{
			NodeName: nodeTypes.GetAbsoluteNodeName(),
			Time:     timestamppb.Now(),
			ResponseType: &recorderpb.RecordResponse_Running{
				Running: &recorderpb.RecordingRunningResponse{
					Stats: &recorderpb.RecordingStatistics{
						BytesCaptured:   stats.BytesWritten,
						PacketsCaptured: stats.PacketsWritten,
						BytesLost:       stats.BytesLost,
						PacketsLost:     stats.PacketsLost,
					}},
			},
		})
		if err != nil {
			// errors are expected if the client disconnects early, therefore
			// we do not log this as an error or warning
			log.WithError(err).Debug("failed to send recording update")
			return
		}

		select {
		case <-h.StatsUpdated:
		case <-h.Done:
			return
		case <-ctx.Done():
			return
		}
	}
}
<commit_msg>hubble/recorder: Fix grpc send from concurrent go routine<commit_after>// Copyright 2021 Authors of Cilium
//
// 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 recorder

import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"os"
	"path"
	"regexp"
	"time"

	recorderpb "github.com/cilium/cilium/api/v1/recorder"
	"github.com/cilium/cilium/pkg/cidr"
	"github.com/cilium/cilium/pkg/hubble/recorder/pcap"
	"github.com/cilium/cilium/pkg/hubble/recorder/recorderoption"
	"github.com/cilium/cilium/pkg/hubble/recorder/sink"
	"github.com/cilium/cilium/pkg/idpool"
	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"
	nodeTypes "github.com/cilium/cilium/pkg/node/types"
	"github.com/cilium/cilium/pkg/recorder"
	"github.com/cilium/cilium/pkg/u8proto"

	"github.com/sirupsen/logrus"
	"google.golang.org/protobuf/types/known/timestamppb"
)

var log = logging.DefaultLogger.WithField(logfields.LogSubsys, "hubble-recorder")

var _ recorderpb.RecorderServer = (*Service)(nil)

const (
	minRuleID = 1
	maxRuleID = 65534

	defaultFileSinkPrefix = "hubble"
)

type Service struct {
	recorder *recorder.Recorder
	dispatch *sink.Dispatch
	ruleIDs  idpool.IDPool
	opts     recorderoption.Options
}

func NewService(r *recorder.Recorder, d *sink.Dispatch, options ...recorderoption.Option) (*Service, error) {
	opts := recorderoption.Default
	for _, o := range options {
		if err := o(&opts); err != nil {
			return nil, err
		}
	}

	if len(opts.StoragePath) == 0 {
		return nil, errors.New("storage path must not be empty")
	}

	if err := os.MkdirAll(opts.StoragePath, 0600); err != nil {
		return nil, fmt.Errorf("failed to create storage path directory: %w", err)
	}

	return &Service{
		recorder: r,
		dispatch: d,
		ruleIDs:  idpool.NewIDPool(minRuleID, maxRuleID),
		opts:     opts,
	}, nil
}

func recordingStoppedResponse(stats sink.Statistics, filePath string) *recorderpb.RecordResponse {
	return &recorderpb.RecordResponse{
		NodeName: nodeTypes.GetAbsoluteNodeName(),
		Time:     timestamppb.Now(),
		ResponseType: &recorderpb.RecordResponse_Stopped{
			Stopped: &recorderpb.RecordingStoppedResponse{
				Stats: &recorderpb.RecordingStatistics{
					BytesCaptured:   stats.BytesWritten,
					PacketsCaptured: stats.PacketsWritten,
					BytesLost:       stats.BytesLost,
					PacketsLost:     stats.PacketsLost,
				},
				Filesink: &recorderpb.FileSinkResult{
					FilePath: filePath,
				},
			},
		},
	}
}

func recordingRunningResponse(stats sink.Statistics) *recorderpb.RecordResponse {
	return &recorderpb.RecordResponse{
		NodeName: nodeTypes.GetAbsoluteNodeName(),
		Time:     timestamppb.Now(),
		ResponseType: &recorderpb.RecordResponse_Running{
			Running: &recorderpb.RecordingRunningResponse{
				Stats: &recorderpb.RecordingStatistics{
					BytesCaptured:   stats.BytesWritten,
					PacketsCaptured: stats.PacketsWritten,
					BytesLost:       stats.BytesLost,
					PacketsLost:     stats.PacketsLost,
				},
			},
		},
	}
}

func (s *Service) Record(stream recorderpb.Recorder_RecordServer) error {
	ctx, cancel := context.WithCancel(stream.Context())
	defer cancel()

	// Spawn a go routine that forwards any received messages in order to be
	// able to use select on it
	reqCh := make(chan *recorderpb.RecordRequest)
	errCh := make(chan error, 1)
	go func() {
		for {
			req, err := stream.Recv()
			if err != nil {
				errCh <- fmt.Errorf("failed to receive from recorder client: %w", err)
				return
			}

			select {
			case reqCh <- req:
			case <-ctx.Done():
				return
			}
		}
	}()

	var (
		recording *sink.Handle
		filePath  string
		err       error
	)

	// Wait for the initial StartRecording message
	select {
	case req := <-reqCh:
		startRecording := req.GetStart()
		if startRecording == nil {
			return fmt.Errorf("received invalid request %q, expected start request", req)
		}

		// The startRecording helper spawns a clean up go routine to remove all
		// state associated with this recording when the context ctx is cancelled.
		recording, filePath, err = s.startRecording(ctx, startRecording)
		if err != nil {
			return err
		}
	case err = <-errCh:
		return err
	case <-ctx.Done():
		return ctx.Err()
	}

	// Send back a confirmation that the recording has started
	err = stream.Send(recordingRunningResponse(recording.Stats()))
	if err != nil {
		return fmt.Errorf("failed to confirmation response: %w", err)
	}

	for {
		select {
		// This case happens when the client has sent us a new request.
		// We expect a start request if recording is nil, and a stop request
		// otherwise.
		case req := <-reqCh:
			if req.GetStop() != nil {
				recording.Stop()
			} else {
				return fmt.Errorf("received invalid request %q, expected stop request", req)
			}
		// This case is hit whenever the recording has updated the statistics (i.e.
		// packets have been captured). We fetch the latest statistics and forward
		// them to the client
		case <-recording.StatsUpdated:
			err = stream.Send(recordingRunningResponse(recording.Stats()))
			if err != nil {
				return fmt.Errorf("failed to send recording running response: %w", err)
			}
		// This case happens when the recording has stopped (i.e. due to the above
		// explicit shutdown or because an error has occurred). If no error has
		// occurred, we assemble the final RecordingStoppedResponse and exit.
		// If an error occurred, we propagate it by returning it from this stub.
		case <-recording.Done:
			err = recording.Err()
			if err != nil {
				return fmt.Errorf("recorder recording error: %w", err)
			}

			err = stream.Send(recordingStoppedResponse(recording.Stats(), filePath))
			if err != nil {
				return fmt.Errorf("failed to send recording stopped response: %w", err)
			}

			return nil
		// The following two cases happen when the client stream is either
		// closed or cancelled. Simply return an error such that it is logged,
		// and exit.
		case err = <-errCh:
			return err
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

const fileExistsRetries = 100

func createPcapFile(basedir, prefix string) (f *os.File, filePath string, err error) {
	try := 0
	for {
		startTime := time.Now().Unix()
		random := rand.Uint32()
		nodeName := nodeTypes.GetAbsoluteNodeName()
		name := fmt.Sprintf("%s_%d_%d_%s.pcap", prefix, startTime, random, nodeName)
		filePath = path.Join(basedir, name)
		f, err = os.OpenFile(filePath, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0600)
		if err != nil {
			if os.IsExist(err) {
				if try++; try < fileExistsRetries {
					continue
				}
			}
			return f, "", fmt.Errorf("failed to create pcap file %q: %w", filePath, err)
		}

		return f, filePath, nil
	}
}

func parseFilters(include []*recorderpb.Filter) ([]recorder.RecorderTuple, error) {
	if len(include) == 0 {
		return nil, errors.New("need to specify at least one include filter")
	}

	filters := []recorder.RecorderTuple{}
	for _, f := range include {
		srcIP, srcPrefix, err := net.ParseCIDR(f.GetSourceCidr())
		if err != nil {
			return nil, fmt.Errorf("failed to parse source cidr %q: %w", f.GetSourceCidr(), err)
		}

		dstIP, dstPrefix, err := net.ParseCIDR(f.GetDestinationCidr())
		if err != nil {
			return nil, fmt.Errorf("failed to parse source cidr %q: %w", f.GetDestinationCidr(), err)
		}

		if (srcIP.To4() == nil) != (dstIP.To4() == nil) {
			return nil, fmt.Errorf("source (%s) and destination cidr (%s) must be same protocol version",
				f.GetSourceCidr(), f.GetDestinationCidr())
		}

		const maxPort = 65535
		if f.GetSourcePort() > maxPort {
			return nil, fmt.Errorf("source port %d out of range", f.GetSourcePort())
		}

		if f.GetDestinationPort() > maxPort {
			return nil, fmt.Errorf("destination port %d out of range", f.GetDestinationPort())
		}

		filters = append(filters, recorder.RecorderTuple{
			SrcPrefix: *cidr.NewCIDR(srcPrefix),
			SrcPort:   uint16(f.GetSourcePort()),
			DstPrefix: *cidr.NewCIDR(dstPrefix),
			DstPort:   uint16(f.GetDestinationPort()),
			Proto:     u8proto.U8proto(f.GetProtocol()),
		})
	}

	return filters, nil
}

var fileSinkPrefixRegex = regexp.MustCompile("^[a-z][a-z0-9]{0,19}$")

// startRecording starts a new recording. It will clean up any state
// associated with the recording if ctx is cancelled or handle.Stop is called.
func (s *Service) startRecording(
	ctx context.Context,
	req *recorderpb.StartRecording,
) (handle *sink.Handle, filePath string, err error) {
	capLen := req.GetMaxCaptureLength()
	prefix := req.GetFilesink().GetFilePrefix()
	if prefix == "" {
		prefix = defaultFileSinkPrefix
	}

	if !fileSinkPrefixRegex.MatchString(prefix) {
		return nil, "", fmt.Errorf("invalid file sink prefix: %q", prefix)
	}

	filters, err := parseFilters(req.GetInclude())
	if err != nil {
		return nil, "", err
	}

	leaseID := s.ruleIDs.LeaseAvailableID()
	ruleID := uint16(leaseID)
	if leaseID == idpool.NoID {
		return nil, "", errors.New("unable to allocate capture rule id")
	}

	var f *os.File
	f, filePath, err = createPcapFile(s.opts.StoragePath, prefix)
	if err != nil {
		return nil, "", err
	}

	defer func() {
		// clean up the recording if any of the subsequent steps fails
		if err != nil {
			_, _ = s.recorder.DeleteRecorder(recorder.ID(ruleID))
			// remove the created pcap file
			_ = f.Close()
			_ = os.Remove(filePath)
			// release will also invalidate the lease
			_ = s.ruleIDs.Release(idpool.ID(ruleID))
		}
	}()

	scopedLog := log.WithFields(logrus.Fields{
		"ruleID":   ruleID,
		"filePath": filePath,
	})
	scopedLog.Debug("starting new recording")

	config := sink.PcapSink{
		RuleID: ruleID,
		Header: pcap.Header{
			SnapshotLength: capLen,
			Datalink:       pcap.Ethernet,
		},
		Writer: pcap.NewWriter(f),
	}

	handle, err = s.dispatch.StartSink(ctx, config)
	if err != nil {
		return nil, "", err
	}

	recInfo := &recorder.RecInfo{
		ID:      recorder.ID(ruleID),
		CapLen:  uint16(capLen),
		Filters: filters,
	}
	_, err = s.recorder.UpsertRecorder(recInfo)
	if err != nil {
		return nil, "", err
	}

	// Ensure to delete the above recorder when the sink has stopped
	go func() {
		<-handle.Done
		scopedLog.Debug("stopping recording")
		_, err := s.recorder.DeleteRecorder(recorder.ID(ruleID))
		if err != nil {
			scopedLog.WithError(err).Warning("failed to delete recorder")
		}
		s.ruleIDs.Release(idpool.ID(ruleID))
	}()

	s.ruleIDs.Use(leaseID)

	return handle, filePath, nil
}
<|endoftext|>


<commit_before>// Copyright 2018-2019 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package kexec

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"reflect"
	"testing"
)

func TestParseMemoryMap(t *testing.T) {
	var mem Memory
	root, err := ioutil.TempDir("", "memmap")
	if err != nil {
		t.Fatalf("Cannot create test dir: %v", err)
	}
	defer os.RemoveAll(root)

	old := memoryMapRoot
	memoryMapRoot = root
	defer func() { memoryMapRoot = old }()

	create := func(dir string, start, end uintptr, typ RangeType) error {
		p := path.Join(root, dir)
		if err := os.Mkdir(p, 0755); err != nil {
			return err
		}
		if err := ioutil.WriteFile(path.Join(p, "start"), []byte(fmt.Sprintf("%#x\n", start)), 0655); err != nil {
			return err
		}
		if err := ioutil.WriteFile(path.Join(p, "end"), []byte(fmt.Sprintf("%#x\n", end)), 0655); err != nil {
			return err
		}
		return ioutil.WriteFile(path.Join(p, "type"), append([]byte(typ), '\n'), 0655)
	}

	if err := create("0", 0, 50, RangeRAM); err != nil {
		t.Fatal(err)
	}
	if err := create("1", 100, 150, RangeNVACPI); err != nil {
		t.Fatal(err)
	}
	if err := create("2", 200, 250, RangeACPI); err != nil {
		t.Fatal(err)
	}
	if err := create("3", 300, 350, RangeNVS); err != nil {
		t.Fatal(err)
	}

	want := []TypedAddressRange{
		{Range: Range{Start: 0, Size: 50}, Type: RangeRAM},
		{Range: Range{Start: 100, Size: 50}, Type: RangeNVACPI},
		{Range: Range{Start: 200, Size: 50}, Type: RangeACPI},
		{Range: Range{Start: 300, Size: 50}, Type: RangeNVS},
	}

	if err := mem.ParseMemoryMap(); err != nil {
		t.Fatalf("ParseMemoryMap() error: %v", err)
	}
	if !reflect.DeepEqual(mem.Phys, want) {
		t.Errorf("ParseMemoryMap() got %v, want %v", mem.Phys, want)
	}
}

func TestAvailableRAM(t *testing.T) {
	old := pageMask
	defer func() {
		pageMask = old
	}()
	// suppose we have 4K pages.
	pageMask = 4095

	var mem Memory
	mem.Phys = []TypedAddressRange{
		TypedAddressRange{Range: Range{Start: 0, Size: 8192}, Type: RangeRAM},
		TypedAddressRange{Range: Range{Start: 8192, Size: 8000}, Type: RangeRAM},
		TypedAddressRange{Range: Range{Start: 20480, Size: 1000}, Type: RangeRAM},
		TypedAddressRange{Range: Range{Start: 24576, Size: 1000}, Type: RangeRAM},
		TypedAddressRange{Range: Range{Start: 28672, Size: 1000}, Type: RangeRAM},
	}

	mem.Segments = []Segment{
		Segment{Phys: Range{Start: 40, Size: 50}},
		Segment{Phys: Range{Start: 8000, Size: 200}},
		Segment{Phys: Range{Start: 18000, Size: 1000}},
		Segment{Phys: Range{Start: 24600, Size: 1000}},
		Segment{Phys: Range{Start: 28000, Size: 10000}},
	}

	want := []TypedAddressRange{
		TypedAddressRange{Range: Range{Start: 0, Size: 40}, Type: RangeRAM},
		TypedAddressRange{Range: Range{Start: 4096, Size: 8000 - 4096}, Type: RangeRAM},
		TypedAddressRange{Range: Range{Start: 12288, Size: 8192 + 8000 - 12288}, Type: RangeRAM},
		TypedAddressRange{Range: Range{Start: 20480, Size: 1000}, Type: RangeRAM},
		TypedAddressRange{Range: Range{Start: 24576, Size: 24}, Type: RangeRAM},
	}

	got := mem.availableRAM()
	if !reflect.DeepEqual(got, want) {
		t.Errorf("availableRAM() got %+v, want %+v", got, want)
	}
}

func TestAlignPhys(t *testing.T) {
	for _, test := range []struct {
		name      string
		seg, want Segment
	}{
		{
			name: "aligned",
			seg: Segment{
				Buf:  Range{Start: 0x1000, Size: 0x1000},
				Phys: Range{Start: 0x2000, Size: 0x1000},
			},
			want: Segment{
				Buf:  Range{Start: 0x1000, Size: 0x1000},
				Phys: Range{Start: 0x2000, Size: 0x1000},
			},
		},
		{
			name: "unaligned",
			seg: Segment{
				Buf:  Range{Start: 0x1011, Size: 0x1022},
				Phys: Range{Start: 0x2011, Size: 0x1022},
			},
			want: Segment{
				Buf:  Range{Start: 0x1000, Size: 0x1033},
				Phys: Range{Start: 0x2000, Size: 0x2000},
			},
		},
		{
			name: "empty_buf",
			seg: Segment{
				Buf:  Range{Start: 0x1011, Size: 0},
				Phys: Range{Start: 0x2011, Size: 0},
			},
			want: Segment{
				Buf:  Range{Start: 0x1000, Size: 0},
				Phys: Range{Start: 0x2000, Size: 0x1000},
			},
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			got := AlignPhys(test.seg)
			if got != test.want {
				t.Errorf("AlignPhys() got %v, want %v", got, test.want)
			}
		})
	}
}

func TestTryMerge(t *testing.T) {
	for _, test := range []struct {
		name   string
		phys   Range
		merged bool
		want   Range
	}{
		{
			name:   "disjunct",
			phys:   Range{Start: 100, Size: 150},
			merged: false,
		},
		{
			name:   "superset",
			phys:   Range{Start: 0, Size: 80},
			merged: true,
			want:   Range{Start: 0, Size: 100},
		},
		{
			name:   "superset",
			phys:   Range{Start: 10, Size: 80},
			merged: true,
			want:   Range{Start: 0, Size: 100},
		},
		{
			name:   "superset",
			phys:   Range{Start: 10, Size: 90},
			merged: true,
			want:   Range{Start: 0, Size: 100},
		},
		{
			name:   "superset",
			phys:   Range{Start: 0, Size: 100},
			merged: true,
			want:   Range{Start: 0, Size: 100},
		},
		{
			name:   "overlap",
			phys:   Range{Start: 0, Size: 150},
			merged: true,
			want:   Range{Start: 0, Size: 150},
		},
		{
			name:   "overlap",
			phys:   Range{Start: 50, Size: 100},
			merged: true,
			want:   Range{Start: 0, Size: 150},
		},
		{
			name:   "overlap",
			phys:   Range{Start: 99, Size: 51},
			merged: true,
			want:   Range{Start: 0, Size: 150},
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			a := NewSegment([]byte("aaaa"), Range{Start: 0, Size: 100})
			b := NewSegment([]byte("bbbb"), test.phys)

			merged := a.tryMerge(b)
			if merged != test.merged {
				t.Fatalf("tryMerge() got %v, want %v", merged, test.merged)
			}
			if !merged {
				return
			}
			if a.Phys != test.want {
				t.Fatalf("Wrong merge result: got %+v, want %+v", a.Phys, test.want)
			}

			got := a.Buf.toSlice()
			want := []byte("aaaabbbb")
			if !bytes.Equal(got, want) {
				t.Errorf("Wrong buf: got %s, want %s", got, want)
			}
		})
	}
}

func TestDedup(t *testing.T) {
	s := []Segment{
		NewSegment([]byte("test"), Range{Start: 0, Size: 100}),
		NewSegment([]byte("test"), Range{Start: 100, Size: 100}),
		NewSegment([]byte("test"), Range{Start: 200, Size: 100}),
		NewSegment([]byte("test"), Range{Start: 250, Size: 50}),
		NewSegment([]byte("test"), Range{Start: 300, Size: 100}),
		NewSegment([]byte("test"), Range{Start: 350, Size: 100}),
	}
	want := []Range{
		Range{Start: 0, Size: 100},
		Range{Start: 100, Size: 100},
		Range{Start: 200, Size: 100},
		Range{Start: 300, Size: 150},
	}

	got := Dedup(s)
	for i := range got {
		if got[i].Phys != want[i] {
			t.Errorf("Dedup() got %v, want %v", got[i].Phys, want[i])
		}
	}

}
<commit_msg>update tests<commit_after>// Copyright 2018-2019 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package kexec

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"reflect"
	"testing"
)

func TestParseMemoryMap(t *testing.T) {
	var mem Memory
	root, err := ioutil.TempDir("", "memmap")
	if err != nil {
		t.Fatalf("Cannot create test dir: %v", err)
	}
	defer os.RemoveAll(root)

	old := memoryMapRoot
	memoryMapRoot = root
	defer func() { memoryMapRoot = old }()

	create := func(dir string, start, end uintptr, typ RangeType) error {
		p := path.Join(root, dir)
		if err := os.Mkdir(p, 0755); err != nil {
			return err
		}
		if err := ioutil.WriteFile(path.Join(p, "start"), []byte(fmt.Sprintf("%#x\n", start)), 0655); err != nil {
			return err
		}
		if err := ioutil.WriteFile(path.Join(p, "end"), []byte(fmt.Sprintf("%#x\n", end)), 0655); err != nil {
			return err
		}
		return ioutil.WriteFile(path.Join(p, "type"), append([]byte(typ), '\n'), 0655)
	}

	if err := create("0", 0, 50, RangeRAM); err != nil {
		t.Fatal(err)
	}
	if err := create("1", 100, 150, RangeACPI); err != nil {
		t.Fatal(err)
	}
	if err := create("2", 200, 250, RangeNVS); err != nil {
		t.Fatal(err)
	}
	if err := create("3", 300, 350, RangeReserved); err != nil {
		t.Fatal(err)
	}

	want := []TypedAddressRange{
		{Range: Range{Start: 0, Size: 50}, Type: RangeRAM},
		{Range: Range{Start: 100, Size: 50}, Type: RangeACPI},
		{Range: Range{Start: 200, Size: 50}, Type: RangeNVS},
		{Range: Range{Start: 300, Size: 50}, Type: RangeReserved},
	}

	if err := mem.ParseMemoryMap(); err != nil {
		t.Fatalf("ParseMemoryMap() error: %v", err)
	}
	if !reflect.DeepEqual(mem.Phys, want) {
		t.Errorf("ParseMemoryMap() got %v, want %v", mem.Phys, want)
	}
}

func TestAvailableRAM(t *testing.T) {
	old := pageMask
	defer func() {
		pageMask = old
	}()
	// suppose we have 4K pages.
	pageMask = 4095

	var mem Memory
	mem.Phys = []TypedAddressRange{
		TypedAddressRange{Range: Range{Start: 0, Size: 8192}, Type: RangeRAM},
		TypedAddressRange{Range: Range{Start: 8192, Size: 8000}, Type: RangeRAM},
		TypedAddressRange{Range: Range{Start: 20480, Size: 1000}, Type: RangeRAM},
		TypedAddressRange{Range: Range{Start: 24576, Size: 1000}, Type: RangeRAM},
		TypedAddressRange{Range: Range{Start: 28672, Size: 1000}, Type: RangeRAM},
	}

	mem.Segments = []Segment{
		Segment{Phys: Range{Start: 40, Size: 50}},
		Segment{Phys: Range{Start: 8000, Size: 200}},
		Segment{Phys: Range{Start: 18000, Size: 1000}},
		Segment{Phys: Range{Start: 24600, Size: 1000}},
		Segment{Phys: Range{Start: 28000, Size: 10000}},
	}

	want := []TypedAddressRange{
		TypedAddressRange{Range: Range{Start: 0, Size: 40}, Type: RangeRAM},
		TypedAddressRange{Range: Range{Start: 4096, Size: 8000 - 4096}, Type: RangeRAM},
		TypedAddressRange{Range: Range{Start: 12288, Size: 8192 + 8000 - 12288}, Type: RangeRAM},
		TypedAddressRange{Range: Range{Start: 20480, Size: 1000}, Type: RangeRAM},
		TypedAddressRange{Range: Range{Start: 24576, Size: 24}, Type: RangeRAM},
	}

	got := mem.availableRAM()
	if !reflect.DeepEqual(got, want) {
		t.Errorf("availableRAM() got %+v, want %+v", got, want)
	}
}

func TestAlignPhys(t *testing.T) {
	for _, test := range []struct {
		name      string
		seg, want Segment
	}{
		{
			name: "aligned",
			seg: Segment{
				Buf:  Range{Start: 0x1000, Size: 0x1000},
				Phys: Range{Start: 0x2000, Size: 0x1000},
			},
			want: Segment{
				Buf:  Range{Start: 0x1000, Size: 0x1000},
				Phys: Range{Start: 0x2000, Size: 0x1000},
			},
		},
		{
			name: "unaligned",
			seg: Segment{
				Buf:  Range{Start: 0x1011, Size: 0x1022},
				Phys: Range{Start: 0x2011, Size: 0x1022},
			},
			want: Segment{
				Buf:  Range{Start: 0x1000, Size: 0x1033},
				Phys: Range{Start: 0x2000, Size: 0x2000},
			},
		},
		{
			name: "empty_buf",
			seg: Segment{
				Buf:  Range{Start: 0x1011, Size: 0},
				Phys: Range{Start: 0x2011, Size: 0},
			},
			want: Segment{
				Buf:  Range{Start: 0x1000, Size: 0},
				Phys: Range{Start: 0x2000, Size: 0x1000},
			},
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			got := AlignPhys(test.seg)
			if got != test.want {
				t.Errorf("AlignPhys() got %v, want %v", got, test.want)
			}
		})
	}
}

func TestTryMerge(t *testing.T) {
	for _, test := range []struct {
		name   string
		phys   Range
		merged bool
		want   Range
	}{
		{
			name:   "disjunct",
			phys:   Range{Start: 100, Size: 150},
			merged: false,
		},
		{
			name:   "superset",
			phys:   Range{Start: 0, Size: 80},
			merged: true,
			want:   Range{Start: 0, Size: 100},
		},
		{
			name:   "superset",
			phys:   Range{Start: 10, Size: 80},
			merged: true,
			want:   Range{Start: 0, Size: 100},
		},
		{
			name:   "superset",
			phys:   Range{Start: 10, Size: 90},
			merged: true,
			want:   Range{Start: 0, Size: 100},
		},
		{
			name:   "superset",
			phys:   Range{Start: 0, Size: 100},
			merged: true,
			want:   Range{Start: 0, Size: 100},
		},
		{
			name:   "overlap",
			phys:   Range{Start: 0, Size: 150},
			merged: true,
			want:   Range{Start: 0, Size: 150},
		},
		{
			name:   "overlap",
			phys:   Range{Start: 50, Size: 100},
			merged: true,
			want:   Range{Start: 0, Size: 150},
		},
		{
			name:   "overlap",
			phys:   Range{Start: 99, Size: 51},
			merged: true,
			want:   Range{Start: 0, Size: 150},
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			a := NewSegment([]byte("aaaa"), Range{Start: 0, Size: 100})
			b := NewSegment([]byte("bbbb"), test.phys)

			merged := a.tryMerge(b)
			if merged != test.merged {
				t.Fatalf("tryMerge() got %v, want %v", merged, test.merged)
			}
			if !merged {
				return
			}
			if a.Phys != test.want {
				t.Fatalf("Wrong merge result: got %+v, want %+v", a.Phys, test.want)
			}

			got := a.Buf.toSlice()
			want := []byte("aaaabbbb")
			if !bytes.Equal(got, want) {
				t.Errorf("Wrong buf: got %s, want %s", got, want)
			}
		})
	}
}

func TestDedup(t *testing.T) {
	s := []Segment{
		NewSegment([]byte("test"), Range{Start: 0, Size: 100}),
		NewSegment([]byte("test"), Range{Start: 100, Size: 100}),
		NewSegment([]byte("test"), Range{Start: 200, Size: 100}),
		NewSegment([]byte("test"), Range{Start: 250, Size: 50}),
		NewSegment([]byte("test"), Range{Start: 300, Size: 100}),
		NewSegment([]byte("test"), Range{Start: 350, Size: 100}),
	}
	want := []Range{
		Range{Start: 0, Size: 100},
		Range{Start: 100, Size: 100},
		Range{Start: 200, Size: 100},
		Range{Start: 300, Size: 150},
	}

	got := Dedup(s)
	for i := range got {
		if got[i].Phys != want[i] {
			t.Errorf("Dedup() got %v, want %v", got[i].Phys, want[i])
		}
	}

}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes 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 kubectl

import (
	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	api "k8s.io/kubernetes/pkg/apis/core"
	"k8s.io/kubernetes/pkg/printers"
)

// FilterFunc is a function that knows how to filter a specific resource kind.
// It receives a generic runtime.Object which must be type-checked by the function.
// Returns a boolean value true if a resource is filtered, or false otherwise.
type FilterFunc func(runtime.Object, printers.PrintOptions) bool

// Filters is a collection of filter funcs
type Filters []FilterFunc

func NewResourceFilter() Filters {
	return []FilterFunc{
		filterPods,
	}
}

// filterPods returns true if a pod should be skipped.
// If show-all is true, the pod will be never be skipped (return false);
// otherwise, skip terminated pod.
func filterPods(obj runtime.Object, options printers.PrintOptions) bool {
	if options.ShowAll {
		return false
	}

	switch p := obj.(type) {
	case *v1.Pod:
		return p.Status.Phase == v1.PodSucceeded || p.Status.Phase == v1.PodFailed
	case *api.Pod:
		return p.Status.Phase == api.PodSucceeded || p.Status.Phase == api.PodFailed
	case *unstructured.Unstructured:
		if (p.GroupVersionKind().GroupKind() != schema.GroupKind{Kind: "Pod"}) {
			return false
		}
		status, ok := p.Object["status"]
		if !ok {
			return false
		}
		statusMap, ok := status.(map[string]interface{})
		if !ok {
			return false
		}
		phase, ok := statusMap["phase"]
		if !ok {
			return false
		}
		phaseString, ok := phase.(string)
		if !ok {
			return false
		}
		return phaseString == string(api.PodSucceeded) || phaseString == string(api.PodFailed)
	}
	return false
}

// Filter loops through a collection of FilterFuncs until it finds one that can filter the given resource
func (f Filters) Filter(obj runtime.Object, opts *printers.PrintOptions) (bool, error) {
	for _, filter := range f {
		if ok := filter(obj, *opts); ok {
			return true, nil
		}
	}
	return false, nil
}
<commit_msg>Revert "fix resource filter for generic printers on get"<commit_after>/*
Copyright 2016 The Kubernetes 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 kubectl

import (
	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/runtime"
	api "k8s.io/kubernetes/pkg/apis/core"
	"k8s.io/kubernetes/pkg/printers"
)

// FilterFunc is a function that knows how to filter a specific resource kind.
// It receives a generic runtime.Object which must be type-checked by the function.
// Returns a boolean value true if a resource is filtered, or false otherwise.
type FilterFunc func(runtime.Object, printers.PrintOptions) bool

// Filters is a collection of filter funcs
type Filters []FilterFunc

func NewResourceFilter() Filters {
	return []FilterFunc{
		filterPods,
	}
}

// filterPods returns true if a pod should be skipped.
// If show-all is true, the pod will be never be skipped (return false);
// otherwise, skip terminated pod.
func filterPods(obj runtime.Object, options printers.PrintOptions) bool {
	if options.ShowAll {
		return false
	}

	switch p := obj.(type) {
	case *v1.Pod:
		return p.Status.Phase == v1.PodSucceeded || p.Status.Phase == v1.PodFailed
	case *api.Pod:
		return p.Status.Phase == api.PodSucceeded || p.Status.Phase == api.PodFailed
	}
	return false
}

// Filter loops through a collection of FilterFuncs until it finds one that can filter the given resource
func (f Filters) Filter(obj runtime.Object, opts *printers.PrintOptions) (bool, error) {
	for _, filter := range f {
		if ok := filter(obj, *opts); ok {
			return true, nil
		}
	}
	return false, nil
}
<|endoftext|>