<commit_before>package ble

import (
	"context"
	"fmt"
	"log"
	"strings"
	"time"

	blelib "github.com/currantlabs/ble"
	"github.com/currantlabs/ble/examples/lib/dev"
	"github.com/pkg/errors"
)

// ClientAdaptor represents a Client Connection to a BLE Peripheral
type ClientAdaptor struct {
	name    string
	address string

	//	uuid    blelib.UUID
	addr    blelib.Addr
	device  blelib.Device
	client  blelib.Client
	profile *blelib.Profile

	connected bool
	ready     chan struct{}
}

// NewClientAdaptor returns a new ClientAdaptor given an address or peripheral name
func NewClientAdaptor(address string) *ClientAdaptor {
	return &ClientAdaptor{
		name:      "BLECLient",
		address:   address,
		connected: false,
	}
}

func (b *ClientAdaptor) Name() string     { return b.name }
func (b *ClientAdaptor) SetName(n string) { b.name = n }
func (b *ClientAdaptor) Address() string  { return b.address }

//func (b *ClientAdaptor) Peripheral() gatt.Peripheral { return b.peripheral }

// Connect initiates a connection to the BLE peripheral. Returns true on successful connection.
func (b *ClientAdaptor) Connect() (err error) {
	d, err := dev.NewDevice("default")
	if err != nil {
		return errors.Wrap(err, "can't new device")
	}
	blelib.SetDefaultDevice(d)
	b.device = d

	var cln blelib.Client

	ctx := blelib.WithSigHandler(context.WithTimeout(context.Background(), 5*time.Second))
	if cln, err = blelib.Connect(ctx, filter(b.Name())); err == nil {
		b.addr = cln.Address()
		b.address = cln.Address().String()
		b.SetName(cln.Name())
	}

	b.client = cln

	p, err := b.client.DiscoverProfile(true)
	if err != nil {
		return err
	}

	b.profile = p
	return
}

// Reconnect attempts to reconnect to the BLE peripheral. If it has an active connection
// it will first close that connection and then establish a new connection.
// Returns true on Successful reconnection
func (b *ClientAdaptor) Reconnect() (err error) {
	if b.connected {
		b.Disconnect()
	}
	return b.Connect()
}

// Disconnect terminates the connection to the BLE peripheral. Returns true on successful disconnect.
func (b *ClientAdaptor) Disconnect() (err error) {
	b.client.CancelConnection()
	return
}

// Finalize finalizes the BLEAdaptor
func (b *ClientAdaptor) Finalize() (err error) {
	return b.Disconnect()
}

// ReadCharacteristic returns bytes from the BLE device for the
// requested characteristic uuid
func (b *ClientAdaptor) ReadCharacteristic(cUUID string) (data []byte, err error) {
	if !b.connected {
		log.Fatalf("Cannot read from BLE device until connected")
		return
	}

	uuid, _ := blelib.Parse(cUUID)

	if u := b.profile.Find(blelib.NewCharacteristic(uuid)); u != nil {
		data, err = b.client.ReadCharacteristic(u.(*blelib.Characteristic))
		if err != nil {
			return nil, err
		}
		fmt.Printf("    Value         %x | %q\n", data, data)
		return data, nil
	}

	return data, nil
}

// WriteCharacteristic writes bytes to the BLE device for the
// requested service and characteristic
func (b *ClientAdaptor) WriteCharacteristic(cUUID string, data []byte) (err error) {
	if !b.connected {
		log.Fatalf("Cannot write to BLE device until connected")
		return
	}

	uuid, _ := blelib.Parse(cUUID)

	if u := b.profile.Find(blelib.NewCharacteristic(uuid)); u != nil {
		err = b.client.WriteCharacteristic(u.(*blelib.Characteristic), data, false)
		if err != nil {
			return err
		}
		return nil
	}

	return
}

// Subscribe subscribes to notifications from the BLE device for the
// requested service and characteristic
func (b *ClientAdaptor) Subscribe(cUUID string, f func([]byte, error)) (err error) {
	if !b.connected {
		log.Fatalf("Cannot subscribe to BLE device until connected")
		return
	}

	uuid, _ := blelib.Parse(cUUID)

	if u := b.profile.Find(blelib.NewCharacteristic(uuid)); u != nil {
		h := func(req []byte) { f(req, nil) }
		err = b.client.Subscribe(u.(*blelib.Characteristic), false, h)
		if err != nil {
			return err
		}
		return nil
	}

	return
}

func filter(name string) blelib.AdvFilter {
	return func(a blelib.Advertisement) bool {
		return strings.ToLower(a.LocalName()) == strings.ToLower(name) ||
			a.Address().String() == strings.ToLower(name)
	}
}
<commit_msg>ble: return error if unable to connect<commit_after>package ble

import (
	"context"
	"fmt"
	"log"
	"strings"

	blelib "github.com/currantlabs/ble"
	"github.com/currantlabs/ble/examples/lib/dev"
	"github.com/pkg/errors"
)

// ClientAdaptor represents a Client Connection to a BLE Peripheral
type ClientAdaptor struct {
	name    string
	address string

	//	uuid    blelib.UUID
	addr    blelib.Addr
	device  blelib.Device
	client  blelib.Client
	profile *blelib.Profile

	connected bool
	ready     chan struct{}
}

// NewClientAdaptor returns a new ClientAdaptor given an address or peripheral name
func NewClientAdaptor(address string) *ClientAdaptor {
	return &ClientAdaptor{
		name:      "BLECLient",
		address:   address,
		connected: false,
	}
}

func (b *ClientAdaptor) Name() string     { return b.name }
func (b *ClientAdaptor) SetName(n string) { b.name = n }
func (b *ClientAdaptor) Address() string  { return b.address }

//func (b *ClientAdaptor) Peripheral() gatt.Peripheral { return b.peripheral }

// Connect initiates a connection to the BLE peripheral. Returns true on successful connection.
func (b *ClientAdaptor) Connect() (err error) {
	d, err := dev.NewDevice("default")
	if err != nil {
		return errors.Wrap(err, "can't new device")
	}
	blelib.SetDefaultDevice(d)
	b.device = d

	var cln blelib.Client

	ctx := blelib.WithSigHandler(context.WithTimeout(context.Background(), 0))
	cln, err = blelib.Connect(ctx, filter(b.Name()))
	if err != nil {
		return errors.Wrap(err, "can't connect")
	}
	b.addr = cln.Address()
	b.address = cln.Address().String()
	b.SetName(cln.Name())
	b.client = cln

	p, err := b.client.DiscoverProfile(true)
	if err != nil {
		return errors.Wrap(err, "can't discover profile")
	}

	b.profile = p
	return
}

// Reconnect attempts to reconnect to the BLE peripheral. If it has an active connection
// it will first close that connection and then establish a new connection.
// Returns true on Successful reconnection
func (b *ClientAdaptor) Reconnect() (err error) {
	if b.connected {
		b.Disconnect()
	}
	return b.Connect()
}

// Disconnect terminates the connection to the BLE peripheral. Returns true on successful disconnect.
func (b *ClientAdaptor) Disconnect() (err error) {
	b.client.CancelConnection()
	return
}

// Finalize finalizes the BLEAdaptor
func (b *ClientAdaptor) Finalize() (err error) {
	return b.Disconnect()
}

// ReadCharacteristic returns bytes from the BLE device for the
// requested characteristic uuid
func (b *ClientAdaptor) ReadCharacteristic(cUUID string) (data []byte, err error) {
	if !b.connected {
		log.Fatalf("Cannot read from BLE device until connected")
		return
	}

	uuid, _ := blelib.Parse(cUUID)

	if u := b.profile.Find(blelib.NewCharacteristic(uuid)); u != nil {
		data, err = b.client.ReadCharacteristic(u.(*blelib.Characteristic))
		if err != nil {
			return nil, err
		}
		fmt.Printf("    Value         %x | %q\n", data, data)
		return data, nil
	}

	return data, nil
}

// WriteCharacteristic writes bytes to the BLE device for the
// requested service and characteristic
func (b *ClientAdaptor) WriteCharacteristic(cUUID string, data []byte) (err error) {
	if !b.connected {
		log.Fatalf("Cannot write to BLE device until connected")
		return
	}

	uuid, _ := blelib.Parse(cUUID)

	if u := b.profile.Find(blelib.NewCharacteristic(uuid)); u != nil {
		err = b.client.WriteCharacteristic(u.(*blelib.Characteristic), data, false)
		if err != nil {
			return err
		}
		return nil
	}

	return
}

// Subscribe subscribes to notifications from the BLE device for the
// requested service and characteristic
func (b *ClientAdaptor) Subscribe(cUUID string, f func([]byte, error)) (err error) {
	if !b.connected {
		log.Fatalf("Cannot subscribe to BLE device until connected")
		return
	}

	uuid, _ := blelib.Parse(cUUID)

	if u := b.profile.Find(blelib.NewCharacteristic(uuid)); u != nil {
		h := func(req []byte) { f(req, nil) }
		err = b.client.Subscribe(u.(*blelib.Characteristic), false, h)
		if err != nil {
			return err
		}
		return nil
	}

	return
}

func filter(name string) blelib.AdvFilter {
	return func(a blelib.Advertisement) bool {
		return strings.ToLower(a.LocalName()) == strings.ToLower(name) ||
			a.Address().String() == strings.ToLower(name)
	}
}
<|endoftext|>


<commit_before>package docker

import (
	"bufio"
	"context"
	"encoding/json"
	"github.com/docker/docker/api/types"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"regexp"
)

const ignoredByteLogSize = 8
const tailSize = `100`

var logWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/logs`)
var eventsWebsocketRequest = regexp.MustCompile(`events`)
var hostCheck = regexp.MustCompile(`vibioh\.fr$`)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return hostCheck.MatchString(r.Host)
	},
}

func upgradeAndAuth(w http.ResponseWriter, r *http.Request) (*websocket.Conn, error) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		defer ws.Close()
		return nil, err
	}

	_, basicAuth, err := ws.ReadMessage()
	if err != nil {
		defer ws.Close()
		return nil, err
	}

	if _, err := isAuthenticatedByBasicAuth(string(basicAuth)); err != nil {
		ws.WriteMessage(websocket.TextMessage, []byte(err.Error()))
		defer ws.Close()
		return nil, err
	}

	return ws, nil
}

func logsContainerWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}

	defer ws.Close()

	logs, err := docker.ContainerLogs(context.Background(), string(containerID), types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true, Follow: true, Tail: tailSize})
	if err != nil {
		log.Print(err)
		return
	}

	defer logs.Close()

	end := make(chan int)

	go func() {
		scanner := bufio.NewScanner(logs)
		for scanner.Scan() {
			select {
			case <-end:
				return
			default:
				logLine := scanner.Bytes()
				if len(logLine) > ignoredByteLogSize {
					if err = ws.WriteMessage(websocket.TextMessage, logLine[ignoredByteLogSize:]); err != nil {
						log.Print(err)
						end <- 1
						return
					}
				}
			}
		}
	}()

	for {
		select {
		case <-end:
			return
		default:
			if _, _, err := ws.NextReader(); err != nil {
				end <- 1
				close(end)
				return
			}
		}
	}
}

func eventsWebsocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, err := upgradeAndAuth(w, r)
	if err != nil {
		return
	}
	defer ws.Close()

	context := context.Background()
	messages, errors := docker.Events(context, types.EventsOptions{})
	defer context.Done()

	end := make(chan int)

	go func() {
		for {
			select {
			case message := <-messages:
				messageJSON, err := json.Marshal(message)
				if err != nil {
					log.Print(err)
					end <- 1
					return
				}

				if err = ws.WriteMessage(websocket.TextMessage, messageJSON); err != nil {
					log.Print(err)
					end <- 1
					return
				}
				break
			case err := <-errors:
				log.Print(err)
				end <- 1
			}
		}
	}()

	for {
		select {
		case <-end:
			return
		default:
			if _, _, err := ws.NextReader(); err != nil {
				end <- 1
				close(end)
				return
			}
		}
	}
}

// WebsocketHandler for Docker Websocket request. Should be use with net/http
type WebsocketHandler struct {
}

func (handler WebsocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	urlPath := []byte(r.URL.Path)

	if logWebsocketRequest.Match(urlPath) {
		logsContainerWebsocketHandler(w, r, logWebsocketRequest.FindSubmatch(urlPath)[1])
	} else if eventsWebsocketRequest.Match((urlPath)) {
		eventsWebsocketHandler(w, r)
	}
}
<commit_msg>Update websocket.go<commit_after>package docker

import (
	"bufio"
	"context"
	"encoding/json"
	"github.com/docker/docker/api/types"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"regexp"
)

const ignoredByteLogSize = 8
const tailSize = `100`

var logWebsocketRequest = regexp.MustCompile(`containers/([^/]+)/logs`)
var eventsWebsocketRequest = regexp.MustCompile(`events`)
var hostCheck = regexp.MustCompile(`vibioh\.fr$`)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return hostCheck.MatchString(r.Host)
	},
}

func upgradeAndAuth(w http.ResponseWriter, r *http.Request) (*websocket.Conn, error) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		defer ws.Close()
		return nil, err
	}

	_, basicAuth, err := ws.ReadMessage()
	if err != nil {
		defer ws.Close()
		return nil, err
	}

	if _, err := isAuthenticatedByBasicAuth(string(basicAuth)); err != nil {
		ws.WriteMessage(websocket.TextMessage, []byte(err.Error()))
		defer ws.Close()
		return nil, err
	}

	return ws, nil
}

func logsContainerWebsocketHandler(w http.ResponseWriter, r *http.Request, containerID []byte) {
	ws, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Print(err)
		return
	}

	defer ws.Close()

	logs, err := docker.ContainerLogs(context.Background(), string(containerID), types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true, Follow: true, Tail: tailSize})
	if err != nil {
		log.Print(err)
		return
	}

	defer logs.Close()

	done := make(chan struct{})

	go func() {
		scanner := bufio.NewScanner(logs)
		for scanner.Scan() {
			select {
			case <-done:
				return

			default:
				logLine := scanner.Bytes()
				if len(logLine) > ignoredByteLogSize && err = ws.WriteMessage(websocket.TextMessage, logLine[ignoredByteLogSize:]); err != nil {
					log.Print(err)
					close(done)
					return
				}
			}
		}
	}()

	for {
		select {
		case <-done:
			return
		default:
			if _, _, err := ws.NextReader(); err != nil {
				close(done)
				return
			}
		}
	}
}

func eventsWebsocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, err := upgradeAndAuth(w, r)
	if err != nil {
		return
	}
	defer ws.Close()

	context := context.Background()
	messages, errors := docker.Events(context, types.EventsOptions{})
	defer context.Done()

	done := make(chan struct{})

	go func() {
		for {
			select {
			case message := <-messages:
				messageJSON, err := json.Marshal(message)
				if err != nil {
					log.Print(err)
					close(done)
					return
				}

				if err = ws.WriteMessage(websocket.TextMessage, messageJSON); err != nil {
					log.Print(err)
					close(done)
					return
				}
				break

			case err := <-errors:
				log.Print(err)
				close(done)
				return

			case <-done:
				return
			}
		}
	}()

	for {
		select {
		case <-done:
			return
		default:
			if _, _, err := ws.NextReader(); err != nil {
				close(done)
				return
			}
		}
	}
}

// WebsocketHandler for Docker Websocket request. Should be use with net/http
type WebsocketHandler struct {
}

func (handler WebsocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	urlPath := []byte(r.URL.Path)

	if logWebsocketRequest.Match(urlPath) {
		logsContainerWebsocketHandler(w, r, logWebsocketRequest.FindSubmatch(urlPath)[1])
	} else if eventsWebsocketRequest.Match((urlPath)) {
		eventsWebsocketHandler(w, r)
	}
}
<|endoftext|>


<commit_before>package docker

import (
	"bufio"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"regexp"
	"strings"

	"github.com/ViBiOh/dashboard/auth"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/gorilla/websocket"
)

const ignoredByteLogSize = 8
const tailSize = `100`
const start = `start`
const stop = `stop`
const busPrefix = `/bus`

var (
	eventsDemand = regexp.MustCompile(`^events (\S+)`)
	logsDemand   = regexp.MustCompile(`^logs (\S+)(?: (.+))?`)
	statsDemand  = regexp.MustCompile(`^stats (\S+)(?: (.+))?`)
)

var (
	eventsPrefix = []byte(`events `)
	logsPrefix   = []byte(`logs `)
	statsPrefix  = []byte(`stats `)
)

var (
	hostCheck       *regexp.Regexp
	websocketOrigin = flag.String(`ws`, `^dashboard`, `Allowed WebSocket Origin pattern`)
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return hostCheck.MatchString(r.Host)
	},
}

// InitWebsocket configure websocket handler
func InitWebsocket() error {
	hostCheck = regexp.MustCompile(*websocketOrigin)

	return nil
}

func upgradeAndAuth(w http.ResponseWriter, r *http.Request) (*websocket.Conn, *auth.User, error) {
	remoteIP := auth.GetRemoteIP(r)
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		defer ws.Close()
		return nil, nil, fmt.Errorf(`Error while upgrading connection: %v`, err)
	}

	_, basicAuth, err := ws.ReadMessage()
	if err != nil {
		defer ws.Close()
		return nil, nil, fmt.Errorf(`Error while reading authentification message: %v`, err)
	}

	user, err := auth.IsAuthenticatedByAuth(string(basicAuth), remoteIP)
	if err != nil {
		ws.WriteMessage(websocket.TextMessage, []byte(err.Error()))
		defer ws.Close()
		return nil, nil, fmt.Errorf(`Error while checking authentification: %v`, err)
	}

	return ws, user, nil
}

func readContent(user *auth.User, ws *websocket.Conn, name string, done chan<- int, content chan<- []byte) {
	for {
		messageType, message, err := ws.ReadMessage()

		if messageType == websocket.CloseMessage {
			close(done)
			return
		}

		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseNoStatusReceived, websocket.CloseAbnormalClosure) {
				log.Printf(`[%s] Error while reading from %s socket: %v`, user.Username, name, err)
			}

			close(done)
			return
		}

		content <- message
	}
}

func streamEvents(ctx context.Context, cancel context.CancelFunc, user *auth.User, _ string, output chan<- []byte) {
	defer cancel()

	filtersArgs := filters.NewArgs()
	labelFilters(user, &filtersArgs, ``)
	eventFilters(&filtersArgs)

	messages, errors := docker.Events(ctx, types.EventsOptions{Filters: filtersArgs})

	for {
		select {
		case <-ctx.Done():
			return

		case message := <-messages:
			messageJSON, err := json.Marshal(message)
			if err != nil {
				log.Printf(`[%s] Events marshalling in error: %v`, user.Username, err)
				cancel()
			} else {
				output <- append(eventsPrefix, messageJSON...)
			}

		case err := <-errors:
			log.Printf(`[%s] Events reading in error: %v`, user.Username, err)
			cancel()
		}
	}
}

func streamLogs(ctx context.Context, cancel context.CancelFunc, user *auth.User, containerID string, output chan<- []byte) {
	logs, err := docker.ContainerLogs(ctx, containerID, types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true, Follow: true, Tail: tailSize})
	defer cancel()

	if err != nil {
		log.Printf(`[%s] Logs opening in error: %v`, user.Username, err)
		return
	}
	defer logs.Close()

	scanner := bufio.NewScanner(logs)
	for scanner.Scan() {
		logLine := scanner.Bytes()
		if len(logLine) > ignoredByteLogSize {
			output <- append(logsPrefix, logLine[ignoredByteLogSize:]...)
		}
	}
}

func streamStats(ctx context.Context, cancel context.CancelFunc, user *auth.User, containerID string, output chan<- []byte) {
	stats, err := docker.ContainerStats(ctx, containerID, true)
	defer cancel()

	if err != nil {
		log.Printf(`[%s] Stats opening in error for %s: %v`, user.Username, containerID, err)
		return
	}
	defer stats.Body.Close()

	scanner := bufio.NewScanner(stats.Body)
	for scanner.Scan() {
		output <- append(statsPrefix, scanner.Bytes()...)
	}
}

func handleBusDemand(user *auth.User, name string, input []byte, demand *regexp.Regexp, cancel context.CancelFunc, output chan<- []byte, streamFn func(context.Context, context.CancelFunc, *auth.User, string, chan<- []byte)) context.CancelFunc {
	demandGroups := demand.FindSubmatch(input)
	if len(demandGroups) < 2 {
		log.Printf(`[%s] Unable to parse bus demand %s for %s`, user.Username, input, name)
	}

	action := string(demandGroups[1])

	containerID := ``
	if len(demandGroups) > 2 {
		containerID = string(demandGroups[2])
	}

	if action == stop && cancel != nil {
		cancel()
	} else if action == start {
		if cancel != nil {
			cancel()
		}

		ctx, newCancel := context.WithCancel(context.Background())
		go streamFn(ctx, newCancel, user, string(containerID), output)

		return newCancel
	}

	return nil
}

func busWebsocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Printf(`Error while upgrading connection to websocket: %v`, err)
		return
	}
	defer ws.Close()

	done := make(chan int)

	output := make(chan []byte)
	defer close(output)

	input := make(chan []byte)
	defer close(input)

	go readContent(user, ws, `streaming`, done, input)

	var eventsCancelFunc context.CancelFunc
	var logsCancelFunc context.CancelFunc
	var statsCancelFunc context.CancelFunc

	if err = ws.WriteMessage(websocket.TextMessage, []byte(`ready`)); err != nil {
		log.Printf(`[%s] Error while saying ready: %v`, user.Username, err)
	}

	for {
		select {
		case <-done:
			return

		case inputBytes := <-input:
			if eventsDemand.Match(inputBytes) {
				eventsCancelFunc = handleBusDemand(user, `events`, inputBytes, eventsDemand, eventsCancelFunc, output, streamEvents)
				if eventsCancelFunc != nil {
					defer eventsCancelFunc()
				}
			} else if logsDemand.Match(inputBytes) {
				logsCancelFunc = handleBusDemand(user, `logs`, inputBytes, logsDemand, logsCancelFunc, output, streamLogs)
				if logsCancelFunc != nil {
					defer logsCancelFunc()
				}
			} else if statsDemand.Match(inputBytes) {
				statsCancelFunc = handleBusDemand(user, `stats`, inputBytes, statsDemand, statsCancelFunc, output, streamStats)
				if statsCancelFunc != nil {
					defer statsCancelFunc()
				}
			}

		case outputBytes := <-output:
			if err = ws.WriteMessage(websocket.TextMessage, outputBytes); err != nil {
				log.Printf(`[%s] Error while writing to streaming: %v`, user.Username, err)
				close(done)
			}
		}
	}
}

// WebsocketHandler for Docker Websocket request. Should be use with net/http
type WebsocketHandler struct {
}

func (handler WebsocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if strings.HasPrefix(r.URL.Path, busPrefix) {
		busWebsocketHandler(w, r)
	}
}
<commit_msg>Adding trace for getting remote ip<commit_after>package docker

import (
	"bufio"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"regexp"
	"strings"

	"github.com/ViBiOh/dashboard/auth"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/gorilla/websocket"
)

const ignoredByteLogSize = 8
const tailSize = `100`
const start = `start`
const stop = `stop`
const busPrefix = `/bus`

var (
	eventsDemand = regexp.MustCompile(`^events (\S+)`)
	logsDemand   = regexp.MustCompile(`^logs (\S+)(?: (.+))?`)
	statsDemand  = regexp.MustCompile(`^stats (\S+)(?: (.+))?`)
)

var (
	eventsPrefix = []byte(`events `)
	logsPrefix   = []byte(`logs `)
	statsPrefix  = []byte(`stats `)
)

var (
	hostCheck       *regexp.Regexp
	websocketOrigin = flag.String(`ws`, `^dashboard`, `Allowed WebSocket Origin pattern`)
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return hostCheck.MatchString(r.Host)
	},
}

// InitWebsocket configure websocket handler
func InitWebsocket() error {
	hostCheck = regexp.MustCompile(*websocketOrigin)

	return nil
}

func upgradeAndAuth(w http.ResponseWriter, r *http.Request) (*websocket.Conn, *auth.User, error) {
	log.Printf(`headers: %v
remoteaddr: %v
request: %v`, r.Header, r.RemoteAddr, r)
	remoteIP := auth.GetRemoteIP(r)
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		defer ws.Close()
		return nil, nil, fmt.Errorf(`Error while upgrading connection: %v`, err)
	}

	_, basicAuth, err := ws.ReadMessage()
	if err != nil {
		defer ws.Close()
		return nil, nil, fmt.Errorf(`Error while reading authentification message: %v`, err)
	}

	user, err := auth.IsAuthenticatedByAuth(string(basicAuth), remoteIP)
	if err != nil {
		ws.WriteMessage(websocket.TextMessage, []byte(err.Error()))
		defer ws.Close()
		return nil, nil, fmt.Errorf(`Error while checking authentification: %v`, err)
	}

	return ws, user, nil
}

func readContent(user *auth.User, ws *websocket.Conn, name string, done chan<- int, content chan<- []byte) {
	for {
		messageType, message, err := ws.ReadMessage()

		if messageType == websocket.CloseMessage {
			close(done)
			return
		}

		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway, websocket.CloseNoStatusReceived, websocket.CloseAbnormalClosure) {
				log.Printf(`[%s] Error while reading from %s socket: %v`, user.Username, name, err)
			}

			close(done)
			return
		}

		content <- message
	}
}

func streamEvents(ctx context.Context, cancel context.CancelFunc, user *auth.User, _ string, output chan<- []byte) {
	defer cancel()

	filtersArgs := filters.NewArgs()
	labelFilters(user, &filtersArgs, ``)
	eventFilters(&filtersArgs)

	messages, errors := docker.Events(ctx, types.EventsOptions{Filters: filtersArgs})

	for {
		select {
		case <-ctx.Done():
			return

		case message := <-messages:
			messageJSON, err := json.Marshal(message)
			if err != nil {
				log.Printf(`[%s] Events marshalling in error: %v`, user.Username, err)
				cancel()
			} else {
				output <- append(eventsPrefix, messageJSON...)
			}

		case err := <-errors:
			log.Printf(`[%s] Events reading in error: %v`, user.Username, err)
			cancel()
		}
	}
}

func streamLogs(ctx context.Context, cancel context.CancelFunc, user *auth.User, containerID string, output chan<- []byte) {
	logs, err := docker.ContainerLogs(ctx, containerID, types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true, Follow: true, Tail: tailSize})
	defer cancel()

	if err != nil {
		log.Printf(`[%s] Logs opening in error: %v`, user.Username, err)
		return
	}
	defer logs.Close()

	scanner := bufio.NewScanner(logs)
	for scanner.Scan() {
		logLine := scanner.Bytes()
		if len(logLine) > ignoredByteLogSize {
			output <- append(logsPrefix, logLine[ignoredByteLogSize:]...)
		}
	}
}

func streamStats(ctx context.Context, cancel context.CancelFunc, user *auth.User, containerID string, output chan<- []byte) {
	stats, err := docker.ContainerStats(ctx, containerID, true)
	defer cancel()

	if err != nil {
		log.Printf(`[%s] Stats opening in error for %s: %v`, user.Username, containerID, err)
		return
	}
	defer stats.Body.Close()

	scanner := bufio.NewScanner(stats.Body)
	for scanner.Scan() {
		output <- append(statsPrefix, scanner.Bytes()...)
	}
}

func handleBusDemand(user *auth.User, name string, input []byte, demand *regexp.Regexp, cancel context.CancelFunc, output chan<- []byte, streamFn func(context.Context, context.CancelFunc, *auth.User, string, chan<- []byte)) context.CancelFunc {
	demandGroups := demand.FindSubmatch(input)
	if len(demandGroups) < 2 {
		log.Printf(`[%s] Unable to parse bus demand %s for %s`, user.Username, input, name)
	}

	action := string(demandGroups[1])

	containerID := ``
	if len(demandGroups) > 2 {
		containerID = string(demandGroups[2])
	}

	if action == stop && cancel != nil {
		cancel()
	} else if action == start {
		if cancel != nil {
			cancel()
		}

		ctx, newCancel := context.WithCancel(context.Background())
		go streamFn(ctx, newCancel, user, string(containerID), output)

		return newCancel
	}

	return nil
}

func busWebsocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, user, err := upgradeAndAuth(w, r)
	if err != nil {
		log.Printf(`Error while upgrading connection to websocket: %v`, err)
		return
	}
	defer ws.Close()

	done := make(chan int)

	output := make(chan []byte)
	defer close(output)

	input := make(chan []byte)
	defer close(input)

	go readContent(user, ws, `streaming`, done, input)

	var eventsCancelFunc context.CancelFunc
	var logsCancelFunc context.CancelFunc
	var statsCancelFunc context.CancelFunc

	if err = ws.WriteMessage(websocket.TextMessage, []byte(`ready`)); err != nil {
		log.Printf(`[%s] Error while saying ready: %v`, user.Username, err)
	}

	for {
		select {
		case <-done:
			return

		case inputBytes := <-input:
			if eventsDemand.Match(inputBytes) {
				eventsCancelFunc = handleBusDemand(user, `events`, inputBytes, eventsDemand, eventsCancelFunc, output, streamEvents)
				if eventsCancelFunc != nil {
					defer eventsCancelFunc()
				}
			} else if logsDemand.Match(inputBytes) {
				logsCancelFunc = handleBusDemand(user, `logs`, inputBytes, logsDemand, logsCancelFunc, output, streamLogs)
				if logsCancelFunc != nil {
					defer logsCancelFunc()
				}
			} else if statsDemand.Match(inputBytes) {
				statsCancelFunc = handleBusDemand(user, `stats`, inputBytes, statsDemand, statsCancelFunc, output, streamStats)
				if statsCancelFunc != nil {
					defer statsCancelFunc()
				}
			}

		case outputBytes := <-output:
			if err = ws.WriteMessage(websocket.TextMessage, outputBytes); err != nil {
				log.Printf(`[%s] Error while writing to streaming: %v`, user.Username, err)
				close(done)
			}
		}
	}
}

// WebsocketHandler for Docker Websocket request. Should be use with net/http
type WebsocketHandler struct {
}

func (handler WebsocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if strings.HasPrefix(r.URL.Path, busPrefix) {
		busWebsocketHandler(w, r)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015-2017 trivago GmbH
//
// 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 awsS3

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/sirupsen/logrus"
	"github.com/trivago/gollum/core/components"
	"strings"
	"sync"
)

const minUploadPartSize = 5 * 1024 * 1024

// BatchedFileWriterInterface extends the components.BatchedWriter interface for rotation checks
type BatchedFileWriterInterface interface {
	components.BatchedWriter
	GetUploadCount() int
}

// BatchedFileWriter is the file producer core.BatchedWriter implementation for the core.BatchedWriterAssembly
type BatchedFileWriter struct {
	s3Client    *s3.S3
	s3Bucket    string
	s3SubFolder string
	fileName    string
	logger      logrus.FieldLogger

	currentMultiPart int64               // current multipart count
	s3UploadID       *string             // upload id from s3 for active file
	totalSize        int                 // total size off all writes to this writer (need for rotations)
	completedParts   []*s3.CompletedPart // collection of uploaded parts

	// need separate byte buffer for min 5mb part uploads.
	// @see http://docs.aws.amazon.com/AmazonS3/latest/API/mpUploadComplete.html
	activeBuffer *s3ByteBuffer

	//todo: check necessity
	writeGuard *sync.Mutex
}

// NewBatchedFileWriter returns a BatchedFileWriter instance
func NewBatchedFileWriter(s3Client *s3.S3, bucket string, fileName string, logger logrus.FieldLogger) BatchedFileWriter {
	var s3Bucket, s3SubFolder string

	if strings.Contains(bucket, "/") {
		split := strings.SplitN(bucket, "/", 2)
		s3Bucket, s3SubFolder = split[0], split[1]
	} else {
		s3Bucket = bucket
		s3SubFolder = ""
	}

	batchedFileWriter := BatchedFileWriter{
		s3Client:    s3Client,
		s3Bucket:    s3Bucket,
		s3SubFolder: s3SubFolder,
		fileName:    fileName,
		logger:      logger,
	}

	batchedFileWriter.init()
	return batchedFileWriter
}

// init BatchedFileWriter struct
func (w *BatchedFileWriter) init() {
	w.totalSize = 0
	w.currentMultiPart = 0
	w.completedParts = []*s3.CompletedPart{}
	w.writeGuard = new(sync.Mutex)
	w.activeBuffer = newS3ByteBuffer()

	w.createMultipartUpload()
}

// Write is part of the BatchedWriter interface and wraps the file.Write() implementation
func (w *BatchedFileWriter) Write(p []byte) (n int, err error) {
	w.activeBuffer.Write(p)

	if size, _ := w.activeBuffer.Size(); size >= minUploadPartSize {
		w.logger.WithField("size", size).Debug("Buffer size ready for request")
		w.uploadPartInput()
	} else {
		w.logger.WithField("size", size).Debug("Buffer size not big enough vor request")
	}

	length := len(p)
	w.totalSize += length
	return length, nil
}

// Name is part of the BatchedWriter interface and wraps the file.Name() implementation
func (w *BatchedFileWriter) Name() string {
	return w.fileName
}

// Size is part of the BatchedWriter interface and wraps the file.Stat().Size() implementation
func (w *BatchedFileWriter) Size() int64 {
	return int64(w.totalSize)
}

// IsAccessible is part of the BatchedWriter interface and check if the writer can access his file
func (w *BatchedFileWriter) IsAccessible() bool {
	return w.s3UploadID != nil
}

// Close is part of the Close interface and handle the file close or compression call
func (w *BatchedFileWriter) Close() error {
	// flush upload buffer
	err := w.uploadPartInput()

	w.completeMultipartUpload()

	return err
}

// GetUploadCount returns the count of completed part uploads
func (w *BatchedFileWriter) GetUploadCount() int {
	return len(w.completedParts)
}

func (w *BatchedFileWriter) getS3Path() string {
	if w.s3SubFolder != "" {
		return fmt.Sprintf("%s/%s", w.s3SubFolder, w.Name())
	}

	return w.Name()
}

func (w *BatchedFileWriter) uploadPartInput() (err error) {
	if size, _ := w.activeBuffer.Size(); size < 1 {
		w.logger.Warning("uploadPartInput(): empty buffer - no upload necessary")
		return
	}

	//todo: check guard necessity
	//w.writeGuard.Lock()
	//defer w.writeGuard.Unlock()

	// increase and get currentMultiPart count
	w.currentMultiPart++
	currentMultiPart := w.currentMultiPart

	// get and reset active buffer
	buffer := w.activeBuffer
	w.activeBuffer = newS3ByteBuffer()
	//w.writeGuard.Unlock()

	input := &s3.UploadPartInput{
		Body:       buffer,
		Bucket:     aws.String(w.s3Bucket),
		Key:        aws.String(w.getS3Path()),
		PartNumber: aws.Int64(currentMultiPart),
		UploadId:   w.s3UploadID,
	}

	result, err := w.s3Client.UploadPart(input)
	if err != nil {
		w.logger.WithError(err).WithField("input", input).Errorf("Can't upload part '%d'", currentMultiPart)
		return
	}

	w.logger.
		WithField("part", currentMultiPart).
		WithField("result", result).
		Debug("upload part successfully send")

	completedPart := s3.CompletedPart{}
	completedPart.SetETag(*result.ETag)
	completedPart.SetPartNumber(currentMultiPart)

	w.completedParts = append(w.completedParts, &completedPart)

	return
}

func (w *BatchedFileWriter) createMultipartUpload() {
	input := &s3.CreateMultipartUploadInput{
		Bucket: aws.String(w.s3Bucket),
		Key:    aws.String(w.getS3Path()),
	}

	result, err := w.s3Client.CreateMultipartUpload(input)
	if err != nil {
		w.logger.WithError(err).WithField("file", w.Name()).Error("Can't create multipart upload")
		return
	}

	w.s3UploadID = result.UploadId
	w.logger.WithField("uploadId", result.UploadId).Debug("successfully created multipart upload")
}

func (w *BatchedFileWriter) completeMultipartUpload() {
	if w.currentMultiPart < 1 {
		w.logger.Warning("No completeMultipartUpload request necessary for zero parts")
		return
	}

	input := &s3.CompleteMultipartUploadInput{
		Bucket:   aws.String(w.s3Bucket),
		Key:      aws.String(w.getS3Path()),
		UploadId: w.s3UploadID,
		MultipartUpload: &s3.CompletedMultipartUpload{
			Parts: w.completedParts,
		},
	}

	result, err := w.s3Client.CompleteMultipartUpload(input)
	if err != nil {
		w.logger.WithError(err).
			WithField("input", input).
			WithField("response", result).
			Error("Can't complete multipart upload")
		return
	}

	w.logger.
		WithField("location", *result.Location).
		WithField("parts", len(w.completedParts)).
		Debug("successfully completed MultipartUpload")
	w.s3UploadID = nil // reset upload id
}
<commit_msg>removed unused sync.Mutex<commit_after>// Copyright 2015-2017 trivago GmbH
//
// 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 awsS3

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/sirupsen/logrus"
	"github.com/trivago/gollum/core/components"
	"strings"
)

const minUploadPartSize = 5 * 1024 * 1024

// BatchedFileWriterInterface extends the components.BatchedWriter interface for rotation checks
type BatchedFileWriterInterface interface {
	components.BatchedWriter
	GetUploadCount() int
}

// BatchedFileWriter is the file producer core.BatchedWriter implementation for the core.BatchedWriterAssembly
type BatchedFileWriter struct {
	s3Client    *s3.S3
	s3Bucket    string
	s3SubFolder string
	fileName    string
	logger      logrus.FieldLogger

	currentMultiPart int64               // current multipart count
	s3UploadID       *string             // upload id from s3 for active file
	totalSize        int                 // total size off all writes to this writer (need for rotations)
	completedParts   []*s3.CompletedPart // collection of uploaded parts

	// need separate byte buffer for min 5mb part uploads.
	// @see http://docs.aws.amazon.com/AmazonS3/latest/API/mpUploadComplete.html
	activeBuffer *s3ByteBuffer
}

// NewBatchedFileWriter returns a BatchedFileWriter instance
func NewBatchedFileWriter(s3Client *s3.S3, bucket string, fileName string, logger logrus.FieldLogger) BatchedFileWriter {
	var s3Bucket, s3SubFolder string

	if strings.Contains(bucket, "/") {
		split := strings.SplitN(bucket, "/", 2)
		s3Bucket, s3SubFolder = split[0], split[1]
	} else {
		s3Bucket = bucket
		s3SubFolder = ""
	}

	batchedFileWriter := BatchedFileWriter{
		s3Client:    s3Client,
		s3Bucket:    s3Bucket,
		s3SubFolder: s3SubFolder,
		fileName:    fileName,
		logger:      logger,
	}

	batchedFileWriter.init()
	return batchedFileWriter
}

// init BatchedFileWriter struct
func (w *BatchedFileWriter) init() {
	w.totalSize = 0
	w.currentMultiPart = 0
	w.completedParts = []*s3.CompletedPart{}
	w.activeBuffer = newS3ByteBuffer()

	w.createMultipartUpload()
}

// Write is part of the BatchedWriter interface and wraps the file.Write() implementation
func (w *BatchedFileWriter) Write(p []byte) (n int, err error) {
	w.activeBuffer.Write(p)

	if size, _ := w.activeBuffer.Size(); size >= minUploadPartSize {
		w.logger.WithField("size", size).Debug("Buffer size ready for request")
		w.uploadPartInput()
	} else {
		w.logger.WithField("size", size).Debug("Buffer size not big enough vor request")
	}

	length := len(p)
	w.totalSize += length
	return length, nil
}

// Name is part of the BatchedWriter interface and wraps the file.Name() implementation
func (w *BatchedFileWriter) Name() string {
	return w.fileName
}

// Size is part of the BatchedWriter interface and wraps the file.Stat().Size() implementation
func (w *BatchedFileWriter) Size() int64 {
	return int64(w.totalSize)
}

// IsAccessible is part of the BatchedWriter interface and check if the writer can access his file
func (w *BatchedFileWriter) IsAccessible() bool {
	return w.s3UploadID != nil
}

// Close is part of the Close interface and handle the file close or compression call
func (w *BatchedFileWriter) Close() error {
	// flush upload buffer
	err := w.uploadPartInput()

	w.completeMultipartUpload()

	return err
}

// GetUploadCount returns the count of completed part uploads
func (w *BatchedFileWriter) GetUploadCount() int {
	return len(w.completedParts)
}

func (w *BatchedFileWriter) getS3Path() string {
	if w.s3SubFolder != "" {
		return fmt.Sprintf("%s/%s", w.s3SubFolder, w.Name())
	}

	return w.Name()
}

func (w *BatchedFileWriter) uploadPartInput() (err error) {
	if size, _ := w.activeBuffer.Size(); size < 1 {
		w.logger.Warning("uploadPartInput(): empty buffer - no upload necessary")
		return
	}

	// increase and get currentMultiPart count
	w.currentMultiPart++
	currentMultiPart := w.currentMultiPart

	// get and reset active buffer
	buffer := w.activeBuffer
	w.activeBuffer = newS3ByteBuffer()

	input := &s3.UploadPartInput{
		Body:       buffer,
		Bucket:     aws.String(w.s3Bucket),
		Key:        aws.String(w.getS3Path()),
		PartNumber: aws.Int64(currentMultiPart),
		UploadId:   w.s3UploadID,
	}

	result, err := w.s3Client.UploadPart(input)
	if err != nil {
		w.logger.WithError(err).WithField("input", input).Errorf("Can't upload part '%d'", currentMultiPart)
		return
	}

	w.logger.
		WithField("part", currentMultiPart).
		WithField("result", result).
		Debug("upload part successfully send")

	completedPart := s3.CompletedPart{}
	completedPart.SetETag(*result.ETag)
	completedPart.SetPartNumber(currentMultiPart)

	w.completedParts = append(w.completedParts, &completedPart)

	return
}

func (w *BatchedFileWriter) createMultipartUpload() {
	input := &s3.CreateMultipartUploadInput{
		Bucket: aws.String(w.s3Bucket),
		Key:    aws.String(w.getS3Path()),
	}

	result, err := w.s3Client.CreateMultipartUpload(input)
	if err != nil {
		w.logger.WithError(err).WithField("file", w.Name()).Error("Can't create multipart upload")
		return
	}

	w.s3UploadID = result.UploadId
	w.logger.WithField("uploadId", result.UploadId).Debug("successfully created multipart upload")
}

func (w *BatchedFileWriter) completeMultipartUpload() {
	if w.currentMultiPart < 1 {
		w.logger.Warning("No completeMultipartUpload request necessary for zero parts")
		return
	}

	input := &s3.CompleteMultipartUploadInput{
		Bucket:   aws.String(w.s3Bucket),
		Key:      aws.String(w.getS3Path()),
		UploadId: w.s3UploadID,
		MultipartUpload: &s3.CompletedMultipartUpload{
			Parts: w.completedParts,
		},
	}

	result, err := w.s3Client.CompleteMultipartUpload(input)
	if err != nil {
		w.logger.WithError(err).
			WithField("input", input).
			WithField("response", result).
			Error("Can't complete multipart upload")
		return
	}

	w.logger.
		WithField("location", *result.Location).
		WithField("parts", len(w.completedParts)).
		Debug("successfully completed MultipartUpload")
	w.s3UploadID = nil // reset upload id
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"sync"
	"time"

	"github.com/spf13/cobra"
	"golang.org/x/sys/unix"

	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/util"
	"github.com/lxc/lxd/lxd/vsock"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/logger"
)

var servers = make(map[string]*http.Server, 2)
var errChan = make(chan error)

type cmdAgent struct {
	global *cmdGlobal
}

func (c *cmdAgent) Command() *cobra.Command {
	cmd := &cobra.Command{}
	cmd.Use = "lxd-agent [--debug]"
	cmd.Short = "LXD virtual machine agent"
	cmd.Long = `Description:
  LXD virtual machine agent

  This daemon is to be run inside virtual machines managed by LXD.
  It will normally be started through init scripts present or injected
  into the virtual machine.
`
	cmd.RunE = c.Run

	return cmd
}

func (c *cmdAgent) Run(cmd *cobra.Command, args []string) error {
	// Setup logger.
	err := logger.InitLogger("", "", c.global.flagLogVerbose, c.global.flagLogDebug, nil)
	if err != nil {
		os.Exit(1)
	}

	logger.Info("Starting")
	defer logger.Info("Stopped")

	// Apply the templated files.
	files, err := templatesApply("files/")
	if err != nil {
		return err
	}

	// Sync the hostname.
	if shared.PathExists("/proc/sys/kernel/hostname") && shared.StringInSlice("/etc/hostname", files) {
		// Open the two files.
		src, err := os.Open("/etc/hostname")
		if err != nil {
			return err
		}

		dst, err := os.Create("/proc/sys/kernel/hostname")
		if err != nil {
			return err
		}

		// Copy the data.
		_, err = io.Copy(dst, src)
		if err != nil {
			return err
		}

		// Close the files.
		_ = src.Close()
		err = dst.Close()
		if err != nil {
			return err
		}
	}

	// Run cloud-init.
	if shared.PathExists("/etc/cloud") && shared.StringInSlice("/var/lib/cloud/seed/nocloud-net/meta-data", files) {
		logger.Info("Seeding cloud-init")

		cloudInitPath := "/run/cloud-init"
		if shared.PathExists(cloudInitPath) {
			logger.Info(fmt.Sprintf("Removing %q", cloudInitPath))
			err = os.RemoveAll(cloudInitPath)
			if err != nil {
				return err
			}
		}

		logger.Info("Rebooting")
		_, _ = shared.RunCommand("reboot")

		// Wait up to 5min for the reboot to actually happen, if it doesn't, then move on to allowing connections.
		time.Sleep(300 * time.Second)
	}

	reconfigureNetworkInterfaces()

	// Load the kernel driver.
	logger.Info("Loading vsock module")
	err = util.LoadModule("vsock")
	if err != nil {
		return fmt.Errorf("Unable to load the vsock kernel module: %w", err)
	}

	// Wait for vsock device to appear.
	for i := 0; i < 5; i++ {
		if !shared.PathExists("/dev/vsock") {
			time.Sleep(1 * time.Second)
		}
	}

	// Mount shares from host.
	c.mountHostShares()

	d := newDaemon(c.global.flagLogDebug, c.global.flagLogVerbose)

	// Start the server.
	err = startHTTPServer(d, c.global.flagLogDebug)
	if err != nil {
		return fmt.Errorf("Failed to start HTTP server: %w", err)
	}

	// Create a cancellation context.
	ctx, cancelFunc := context.WithCancel(context.Background())

	// Start status notifier in background.
	cancelStatusNotifier := c.startStatusNotifier(ctx, d.chConnected)

	// Check context ID periodically, and restart the HTTP server if needed.
	go func() {
		for range time.Tick(30 * time.Second) {
			cid, err := vsock.ContextID()
			if err != nil {
				continue
			}

			if d.localCID == cid {
				continue
			}

			// Restart server
			servers["http"].Close()

			err = startHTTPServer(d, c.global.flagLogDebug)
			if err != nil {
				errChan <- err
			}

			// Update context ID.
			d.localCID = cid
		}
	}()

	// Done with early setup, tell systemd to continue boot.
	// Allows a service that needs a file that's generated by the agent to be able to declare After=lxd-agent
	// and know the file will have been created by the time the service is started.
	if os.Getenv("NOTIFY_SOCKET") != "" {
		_, err := shared.RunCommand("systemd-notify", "READY=1")
		if err != nil {
			cancelStatusNotifier() // Ensure STOPPED status is written to QEMU status ringbuffer.
			cancelFunc()

			return fmt.Errorf("Failed to notify systemd of readiness: %w", err)
		}
	}

	// Cancel context when SIGTEM is received.
	chSignal := make(chan os.Signal, 1)
	signal.Notify(chSignal, unix.SIGTERM)

	exitStatus := 0

	select {
	case <-chSignal:
	case err := <-errChan:
		fmt.Fprintln(os.Stderr, err)
		exitStatus = 1
	}

	cancelStatusNotifier() // Ensure STOPPED status is written to QEMU status ringbuffer.
	cancelFunc()

	os.Exit(exitStatus)

	return nil
}

// startStatusNotifier sends status of agent to vserial ring buffer every 10s or when context is done.
// Returns a function that can be used to update the running status to STOPPED in the ring buffer.
func (c *cmdAgent) startStatusNotifier(ctx context.Context, chConnected <-chan struct{}) context.CancelFunc {
	// Write initial started status.
	_ = c.writeStatus("STARTED")

	wg := sync.WaitGroup{}
	exitCtx, exit := context.WithCancel(ctx) // Allows manual synchronous cancellation via cancel function.
	cancel := func() {
		exit()    // Signal for the go routine to end.
		wg.Wait() // Wait for the go routine to actually finish.
	}

	wg.Add(1)
	go func() {
		defer wg.Done() // Signal to cancel function that we are done.

		ticker := time.NewTicker(time.Duration(time.Second) * 5)
		defer ticker.Stop()

		for {
			select {
			case <-chConnected:
				_ = c.writeStatus("CONNECTED") // Indicate we were able to connect to LXD.
			case <-ticker.C:
				_ = c.writeStatus("STARTED") // Re-populate status periodically in case LXD restarts.
			case <-exitCtx.Done():
				_ = c.writeStatus("STOPPED") // Indicate we are stopping to LXD and exit go routine.
				return
			}
		}
	}()

	return cancel
}

// writeStatus writes a status code to the vserial ring buffer used to detect agent status on host.
func (c *cmdAgent) writeStatus(status string) error {
	if shared.PathExists("/dev/virtio-ports/org.linuxcontainers.lxd") {
		vSerial, err := os.OpenFile("/dev/virtio-ports/org.linuxcontainers.lxd", os.O_RDWR, 0600)
		if err != nil {
			return err
		}

		defer vSerial.Close()

		_, err = vSerial.Write([]byte(fmt.Sprintf("%s\n", status)))
		if err != nil {
			return err
		}
	}

	return nil
}

// mountHostShares reads the agent-mounts.json file from config share and mounts the shares requested.
func (c *cmdAgent) mountHostShares() {
	agentMountsFile := "./agent-mounts.json"
	if !shared.PathExists(agentMountsFile) {
		return
	}

	b, err := os.ReadFile(agentMountsFile)
	if err != nil {
		logger.Errorf("Failed to load agent mounts file %q: %v", agentMountsFile, err)
	}

	var agentMounts []instancetype.VMAgentMount
	err = json.Unmarshal(b, &agentMounts)
	if err != nil {
		logger.Errorf("Failed to parse agent mounts file %q: %v", agentMountsFile, err)
		return
	}

	for _, mount := range agentMounts {
		// Convert relative mounts to absolute from / otherwise dir creation fails or mount fails.
		if !strings.HasPrefix(mount.Target, "/") {
			mount.Target = fmt.Sprintf("/%s", mount.Target)
		}

		if !shared.PathExists(mount.Target) {
			err := os.MkdirAll(mount.Target, 0755)
			if err != nil {
				logger.Errorf("Failed to create mount target %q", mount.Target)
				continue // Don't try to mount if mount point can't be created.
			}
		}

		if mount.FSType == "9p" {
			// Before mounting with 9p, try virtio-fs and use 9p as the fallback.
			args := []string{"-t", "virtiofs", mount.Source, mount.Target}

			for _, opt := range mount.Options {
				// Ignore the transport and msize mount option as they are specific to 9p.
				if strings.HasPrefix(opt, "trans=") || strings.HasPrefix(opt, "msize=") {
					continue
				}

				args = append(args, "-o", opt)
			}

			_, err = shared.RunCommand("mount", args...)
			if err == nil {
				logger.Infof("Mounted %q (Type: %q, Options: %v) to %q", mount.Source, "virtiofs", mount.Options, mount.Target)
				continue
			}
		}

		args := []string{"-t", mount.FSType, mount.Source, mount.Target}

		for _, opt := range mount.Options {
			args = append(args, "-o", opt)
		}

		_, err = shared.RunCommand("mount", args...)
		if err != nil {
			logger.Errorf("Failed mount %q (Type: %q, Options: %v) to %q: %v", mount.Source, mount.FSType, mount.Options, mount.Target, err)
			continue
		}

		logger.Infof("Mounted %q (Type: %q, Options: %v) to %q", mount.Source, mount.FSType, mount.Options, mount.Target)
	}
}
<commit_msg>lxd-agent: Start devlxd before calling systemd-notify<commit_after>package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"sync"
	"time"

	"github.com/spf13/cobra"
	"golang.org/x/sys/unix"

	"github.com/lxc/lxd/lxd/instance/instancetype"
	"github.com/lxc/lxd/lxd/util"
	"github.com/lxc/lxd/lxd/vsock"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/logger"
)

var servers = make(map[string]*http.Server, 2)
var errChan = make(chan error)

type cmdAgent struct {
	global *cmdGlobal
}

func (c *cmdAgent) Command() *cobra.Command {
	cmd := &cobra.Command{}
	cmd.Use = "lxd-agent [--debug]"
	cmd.Short = "LXD virtual machine agent"
	cmd.Long = `Description:
  LXD virtual machine agent

  This daemon is to be run inside virtual machines managed by LXD.
  It will normally be started through init scripts present or injected
  into the virtual machine.
`
	cmd.RunE = c.Run

	return cmd
}

func (c *cmdAgent) Run(cmd *cobra.Command, args []string) error {
	// Setup logger.
	err := logger.InitLogger("", "", c.global.flagLogVerbose, c.global.flagLogDebug, nil)
	if err != nil {
		os.Exit(1)
	}

	logger.Info("Starting")
	defer logger.Info("Stopped")

	// Apply the templated files.
	files, err := templatesApply("files/")
	if err != nil {
		return err
	}

	// Sync the hostname.
	if shared.PathExists("/proc/sys/kernel/hostname") && shared.StringInSlice("/etc/hostname", files) {
		// Open the two files.
		src, err := os.Open("/etc/hostname")
		if err != nil {
			return err
		}

		dst, err := os.Create("/proc/sys/kernel/hostname")
		if err != nil {
			return err
		}

		// Copy the data.
		_, err = io.Copy(dst, src)
		if err != nil {
			return err
		}

		// Close the files.
		_ = src.Close()
		err = dst.Close()
		if err != nil {
			return err
		}
	}

	// Run cloud-init.
	if shared.PathExists("/etc/cloud") && shared.StringInSlice("/var/lib/cloud/seed/nocloud-net/meta-data", files) {
		logger.Info("Seeding cloud-init")

		cloudInitPath := "/run/cloud-init"
		if shared.PathExists(cloudInitPath) {
			logger.Info(fmt.Sprintf("Removing %q", cloudInitPath))
			err = os.RemoveAll(cloudInitPath)
			if err != nil {
				return err
			}
		}

		logger.Info("Rebooting")
		_, _ = shared.RunCommand("reboot")

		// Wait up to 5min for the reboot to actually happen, if it doesn't, then move on to allowing connections.
		time.Sleep(300 * time.Second)
	}

	reconfigureNetworkInterfaces()

	// Load the kernel driver.
	logger.Info("Loading vsock module")
	err = util.LoadModule("vsock")
	if err != nil {
		return fmt.Errorf("Unable to load the vsock kernel module: %w", err)
	}

	// Wait for vsock device to appear.
	for i := 0; i < 5; i++ {
		if !shared.PathExists("/dev/vsock") {
			time.Sleep(1 * time.Second)
		}
	}

	// Mount shares from host.
	c.mountHostShares()

	d := newDaemon(c.global.flagLogDebug, c.global.flagLogVerbose)

	// Start the server.
	err = startHTTPServer(d, c.global.flagLogDebug)
	if err != nil {
		return fmt.Errorf("Failed to start HTTP server: %w", err)
	}

	// Check context ID periodically, and restart the HTTP server if needed.
	go func() {
		for range time.Tick(30 * time.Second) {
			cid, err := vsock.ContextID()
			if err != nil {
				continue
			}

			if d.localCID == cid {
				continue
			}

			// Restart server
			servers["http"].Close()

			err = startHTTPServer(d, c.global.flagLogDebug)
			if err != nil {
				errChan <- err
			}

			// Update context ID.
			d.localCID = cid
		}
	}()

	// Check whether we should start the devlxd server in the early setup. This way, /dev/lxd/sock
	// will be available for any systemd services starting after the lxd-agent.
	if shared.PathExists("agent.conf") {
		f, err := os.Open("agent.conf")
		if err != nil {
			return err
		}

		err = setConnectionInfo(d, f)
		if err != nil {
			_ = f.Close()
			return err
		}

		_ = f.Close()

		if d.devlxdEnabled {
			err = startDevlxdServer(d)
			if err != nil {
				return err
			}
		}
	}

	// Create a cancellation context.
	ctx, cancelFunc := context.WithCancel(context.Background())

	// Start status notifier in background.
	cancelStatusNotifier := c.startStatusNotifier(ctx, d.chConnected)

	// Done with early setup, tell systemd to continue boot.
	// Allows a service that needs a file that's generated by the agent to be able to declare After=lxd-agent
	// and know the file will have been created by the time the service is started.
	if os.Getenv("NOTIFY_SOCKET") != "" {
		_, err := shared.RunCommand("systemd-notify", "READY=1")
		if err != nil {
			cancelStatusNotifier() // Ensure STOPPED status is written to QEMU status ringbuffer.
			cancelFunc()

			return fmt.Errorf("Failed to notify systemd of readiness: %w", err)
		}
	}

	// Cancel context when SIGTEM is received.
	chSignal := make(chan os.Signal, 1)
	signal.Notify(chSignal, unix.SIGTERM)

	exitStatus := 0

	select {
	case <-chSignal:
	case err := <-errChan:
		fmt.Fprintln(os.Stderr, err)
		exitStatus = 1
	}

	cancelStatusNotifier() // Ensure STOPPED status is written to QEMU status ringbuffer.
	cancelFunc()

	os.Exit(exitStatus)

	return nil
}

// startStatusNotifier sends status of agent to vserial ring buffer every 10s or when context is done.
// Returns a function that can be used to update the running status to STOPPED in the ring buffer.
func (c *cmdAgent) startStatusNotifier(ctx context.Context, chConnected <-chan struct{}) context.CancelFunc {
	// Write initial started status.
	_ = c.writeStatus("STARTED")

	wg := sync.WaitGroup{}
	exitCtx, exit := context.WithCancel(ctx) // Allows manual synchronous cancellation via cancel function.
	cancel := func() {
		exit()    // Signal for the go routine to end.
		wg.Wait() // Wait for the go routine to actually finish.
	}

	wg.Add(1)
	go func() {
		defer wg.Done() // Signal to cancel function that we are done.

		ticker := time.NewTicker(time.Duration(time.Second) * 5)
		defer ticker.Stop()

		for {
			select {
			case <-chConnected:
				_ = c.writeStatus("CONNECTED") // Indicate we were able to connect to LXD.
			case <-ticker.C:
				_ = c.writeStatus("STARTED") // Re-populate status periodically in case LXD restarts.
			case <-exitCtx.Done():
				_ = c.writeStatus("STOPPED") // Indicate we are stopping to LXD and exit go routine.
				return
			}
		}
	}()

	return cancel
}

// writeStatus writes a status code to the vserial ring buffer used to detect agent status on host.
func (c *cmdAgent) writeStatus(status string) error {
	if shared.PathExists("/dev/virtio-ports/org.linuxcontainers.lxd") {
		vSerial, err := os.OpenFile("/dev/virtio-ports/org.linuxcontainers.lxd", os.O_RDWR, 0600)
		if err != nil {
			return err
		}

		defer vSerial.Close()

		_, err = vSerial.Write([]byte(fmt.Sprintf("%s\n", status)))
		if err != nil {
			return err
		}
	}

	return nil
}

// mountHostShares reads the agent-mounts.json file from config share and mounts the shares requested.
func (c *cmdAgent) mountHostShares() {
	agentMountsFile := "./agent-mounts.json"
	if !shared.PathExists(agentMountsFile) {
		return
	}

	b, err := os.ReadFile(agentMountsFile)
	if err != nil {
		logger.Errorf("Failed to load agent mounts file %q: %v", agentMountsFile, err)
	}

	var agentMounts []instancetype.VMAgentMount
	err = json.Unmarshal(b, &agentMounts)
	if err != nil {
		logger.Errorf("Failed to parse agent mounts file %q: %v", agentMountsFile, err)
		return
	}

	for _, mount := range agentMounts {
		// Convert relative mounts to absolute from / otherwise dir creation fails or mount fails.
		if !strings.HasPrefix(mount.Target, "/") {
			mount.Target = fmt.Sprintf("/%s", mount.Target)
		}

		if !shared.PathExists(mount.Target) {
			err := os.MkdirAll(mount.Target, 0755)
			if err != nil {
				logger.Errorf("Failed to create mount target %q", mount.Target)
				continue // Don't try to mount if mount point can't be created.
			}
		}

		if mount.FSType == "9p" {
			// Before mounting with 9p, try virtio-fs and use 9p as the fallback.
			args := []string{"-t", "virtiofs", mount.Source, mount.Target}

			for _, opt := range mount.Options {
				// Ignore the transport and msize mount option as they are specific to 9p.
				if strings.HasPrefix(opt, "trans=") || strings.HasPrefix(opt, "msize=") {
					continue
				}

				args = append(args, "-o", opt)
			}

			_, err = shared.RunCommand("mount", args...)
			if err == nil {
				logger.Infof("Mounted %q (Type: %q, Options: %v) to %q", mount.Source, "virtiofs", mount.Options, mount.Target)
				continue
			}
		}

		args := []string{"-t", mount.FSType, mount.Source, mount.Target}

		for _, opt := range mount.Options {
			args = append(args, "-o", opt)
		}

		_, err = shared.RunCommand("mount", args...)
		if err != nil {
			logger.Errorf("Failed mount %q (Type: %q, Options: %v) to %q: %v", mount.Source, mount.FSType, mount.Options, mount.Target, err)
			continue
		}

		logger.Infof("Mounted %q (Type: %q, Options: %v) to %q", mount.Source, mount.FSType, mount.Options, mount.Target)
	}
}
<|endoftext|>


<commit_before>package jobserver

import (
	"errors"
	"github.com/dmaze/goordinate/cborrpc"
	"github.com/dmaze/goordinate/coordinate"
	"github.com/mitchellh/mapstructure"
	"math"
	"reflect"
)

// AddWorkUnits adds any number of work units to a work spec.  Each oy
// the work units is a cborrpc.PythonTuple or slice containing a
// string with the work unit key, a dictionary with the work unit
// data, and an optional dictionary with additional metadata.
func (jobs *JobServer) AddWorkUnits(workSpecName string, workUnitKvp []interface{}) (bool, string, error) {
	spec, err := jobs.Namespace.WorkSpec(workSpecName)
	if err != nil {
		return false, "", err
	}

	// Unmarshal the work unit list into a []AddWorkUnitItem.
	// Fail now if any are invalid.
	items := make([]coordinate.AddWorkUnitItem, len(workUnitKvp))
	for i, kvp := range workUnitKvp {
		items[i], err = coordinate.ExtractAddWorkUnitItem(kvp)
		if err != nil {
			return false, "", err
		}
	}

	// Now go through and add them all
	for _, item := range items {
		_, err = spec.AddWorkUnit(item.Key, item.Data, item.Priority)
		if err != nil {
			// Again, Python coordinate expects to never see
			// a failure here?
			return false, "", err
		}
	}
	return true, "", nil
}

// GetWorkUnitsOptions contains unmarshaled options for GetWorkUnits().
type GetWorkUnitsOptions struct {
	// WorkUnitKeys contains a list of work unit keys to retrieve.
	// If this option is supplied, all other options are ignored.
	WorkUnitKeys []string `mapstructure:"work_unit_keys"`

	// State provides a list of states to query on.  If this is
	// provided then only work units in one of the specified states
	// will be returned.
	State []WorkUnitStatus

	// Start gives a starting point to iterate through the list of
	// work units.  It is the name of the last work unit returned
	// in the previous call to GetWorkUnits().  No work unit whose
	// name is lexicographically less than this will be returned.
	Start string

	// Limit specifies the maximum number of work units to return.
	// Defaults to 1000.
	Limit int
}

// gwuStateHook is a mapstructure decode hook that expands a single int
// or a PythonTuple into a slice of int (WorkUnitStatus).
func gwuStateHook(from reflect.Type, to reflect.Type, data interface{}) (interface{}, error) {
	// to must be []WorkUnitStatus
	if to.Kind() != reflect.Slice || to.Elem().Name() != "WorkUnitStatus" {
		return data, nil
	}
	switch value := data.(type) {
	case cborrpc.PythonTuple:
		// If from is a tuple, return its contents
		return value.Items, nil
	case WorkUnitStatus:
		// Package it into a slice
		return []WorkUnitStatus{value}, nil
	case int:
		// If from is an int, box it
		return []WorkUnitStatus{WorkUnitStatus(value)}, nil
	case uint64:
		return []WorkUnitStatus{WorkUnitStatus(value)}, nil
	default:
		// Otherwise, hope we can deal normally
		return data, nil
	}
}

// GetWorkUnits retrieves the keys and data dictionaries for some number
// of work units.  If options contains "work_unit_keys", those specific
// work units are retrieved; otherwise the work units are based on
// which of GetWorkUnitsOptions are present.
//
// On success, the return value is a slice of cborrpc.PythonTuple
// objects where each contains the work unit key as a byte slice and
// the data dictionary.
func (jobs *JobServer) GetWorkUnits(workSpecName string, options map[string]interface{}) ([]interface{}, string, error) {
	var workUnits map[string]coordinate.WorkUnit
	gwuOptions := GetWorkUnitsOptions{
		Limit: 1000,
	}

	spec, err := jobs.Namespace.WorkSpec(workSpecName)
	var decoder *mapstructure.Decoder
	if err == nil {
		config := mapstructure.DecoderConfig{
			DecodeHook: mapstructure.ComposeDecodeHookFunc(gwuStateHook, cborrpc.DecodeBytesAsString),
			Result:     &gwuOptions,
		}
		decoder, err = mapstructure.NewDecoder(&config)
	}
	if err == nil {
		err = decoder.Decode(options)
	}
	if err == nil {
		query := coordinate.WorkUnitQuery{
			Names: gwuOptions.WorkUnitKeys,
		}
		if gwuOptions.WorkUnitKeys == nil {
			query.PreviousName = gwuOptions.Start
			query.Limit = gwuOptions.Limit
		}
		if gwuOptions.WorkUnitKeys == nil && gwuOptions.State != nil {
			query.Statuses = make([]coordinate.WorkUnitStatus, len(gwuOptions.State))
			for i, state := range gwuOptions.State {
				query.Statuses[i], err = translateWorkUnitStatus(state)
				if err != nil {
					break
				}
			}
		}
		if err == nil {
			workUnits, err = spec.WorkUnits(query)
		}
	}
	if err != nil {
		return nil, "", err
	}
	// The marshalled result is a list of pairs of (key, data).
	var result []interface{}
	for name, unit := range workUnits {
		var data map[string]interface{}
		attempt, err := unit.ActiveAttempt()
		if err == nil && attempt != nil {
			data, err = attempt.Data()
		}
		if err == nil && data == nil {
			data, err = unit.Data()
		}
		if err != nil {
			return nil, "", err
		}
		tuple := cborrpc.PythonTuple{Items: []interface{}{[]byte(name), data}}
		result = append(result, tuple)
	}
	return result, "", nil
}

// GetWorkUnitStatus returns a summary status of zero or more work
// units in a single work spec.  On success, the returned list of
// dictionaries corresponds one-to-one with workUnitKeys.  If there is
// no such work unit, nil is in the list; otherwise each map contains
// keys "status", "expiration", "worker_id", and "traceback".
func (jobs *JobServer) GetWorkUnitStatus(workSpecName string, workUnitKeys []string) ([]map[string]interface{}, string, error) {
	spec, err := jobs.Namespace.WorkSpec(workSpecName)
	if err != nil {
		return nil, "", err
	}

	result := make([]map[string]interface{}, len(workUnitKeys))
	for i, key := range workUnitKeys {
		workUnit, err := spec.WorkUnit(key)
		if err != nil {
			return nil, "", err
		} else if workUnit == nil {
			result[i] = nil
		} else {
			r := make(map[string]interface{})
			status, attempt, err := workUnitStatus(workUnit)
			if err != nil {
				return nil, "", err
			}
			r["status"] = status
			if attempt != nil {
				r["worker_id"] = attempt.Worker().Name()
			}
			if status == Pending && attempt != nil {
				expiration, err := attempt.ExpirationTime()
				if err != nil {
					return nil, "", err
				}
				r["expiration"] = expiration.Unix()
			}
			if status == Failed && attempt != nil {
				data, err := attempt.Data()
				if err != nil {
					return nil, "", err
				}
				if traceback := data["traceback"]; traceback != nil {
					r["traceback"] = traceback
				}
			}
			result[i] = r
		}
	}
	return result, "", nil
}

// PrioritizeWorkUnitsOptions specifies which work units PrioritizeWorkUnits
// should adjust and how.
type PrioritizeWorkUnitsOptions struct {
	// WorkUnitKeys gives the names of the work units to reprioritize.
	// If not present, does nothing.
	WorkUnitKeys []string `mapstructure:"work_unit_keys"`

	// Priority sets an absolute priority.  If a NaN value, make a
	// change specified by Adjustment instead.
	Priority float64

	// Adjustment is added to the priorities of each of the work
	// units, if Priority is NaN.  If also a NaN value, do nothing.
	Adjustment float64
}

// PrioritizeWorkUnits changes the priorities of some number of work
// units.  The actual work units are in options["work_unit_keys"].  A
// higher priority results in the work units being scheduled sooner.
func (jobs *JobServer) PrioritizeWorkUnits(workSpecName string, options map[string]interface{}) (bool, string, error) {
	var (
		err      error
		query    coordinate.WorkUnitQuery
		workSpec coordinate.WorkSpec
	)
	pwuOptions := PrioritizeWorkUnitsOptions{
		Priority:   math.NaN(),
		Adjustment: math.NaN(),
	}
	workSpec, err = jobs.Namespace.WorkSpec(workSpecName)
	if err == nil {
		err = decode(&pwuOptions, options)
	}
	if err == nil && pwuOptions.WorkUnitKeys == nil {
		return false, "missing work_unit_keys", err
	}
	if err == nil {
		query.Names = pwuOptions.WorkUnitKeys
		if !math.IsNaN(pwuOptions.Priority) {
			err = workSpec.SetWorkUnitPriorities(query, pwuOptions.Priority)
		} else if !math.IsNaN(pwuOptions.Adjustment) {
			err = workSpec.AdjustWorkUnitPriorities(query, pwuOptions.Adjustment)
		}
	}
	return err == nil, "", err
}

// CountWorkUnits returns the number of work units in each status for
// a given work spec.
func (jobs *JobServer) CountWorkUnits(workSpecName string) (map[WorkUnitStatus]int, string, error) {
	workSpec, err := jobs.Namespace.WorkSpec(workSpecName)
	if err != nil {
		return nil, "", err
	}

	// TODO(dmaze): This is a bad way to do this; it should be
	// boiled down to a single call in the API

	result := make(map[WorkUnitStatus]int)
	var workUnits map[string]coordinate.WorkUnit
	var prev string
	for {
		workUnits, err = workSpec.WorkUnits(coordinate.WorkUnitQuery{
			PreviousName: prev,
			Limit:        1000,
		})
		if err != nil {
			return nil, "", err
		}
		for name, workUnit := range workUnits {
			status, _, err := workUnitStatus(workUnit)
			if err != nil {
				return nil, "", err
			}
			result[status]++
			if name > prev {
				prev = name
			}
		}
		if len(workUnits) == 0 {
			break
		}
	}
	return result, "", nil
}

// workUnitStatus extracts a summary of the status of a single work
// unit.  This produces its external coordinate status and the active
// attempt (if any) on success.
func workUnitStatus(workUnit coordinate.WorkUnit) (status WorkUnitStatus, attempt coordinate.Attempt, err error) {
	var attemptStatus coordinate.AttemptStatus
	attempt, err = workUnit.ActiveAttempt()
	if err == nil && attempt == nil {
		status = Available
		return
	}
	if err == nil {
		attemptStatus, err = attempt.Status()
	}
	if err == nil {
		switch attemptStatus {
		case coordinate.Pending:
			status = Pending
		case coordinate.Expired:
			status = Available
			attempt = nil
		case coordinate.Finished:
			status = Finished
		case coordinate.Failed:
			status = Failed
		case coordinate.Retryable:
			status = Available
			attempt = nil
		default:
			err = errors.New("unexpected attempt status")
		}
	}
	return
}

// DelWorkUnitsOptions specifies the options for DelWorkUnits.  The
// first of All, WorkUnitKeys, or State given defines the operation to
// perform.  If none of these are given, the zero value for this
// structure tells DelWorkUnits to do nothing.
type DelWorkUnitsOptions struct {
	// All, if set to true, directs DelWorkUnits to delete
	// all of the work units in its work spec.  If this is
	// provided, all other options are ignored.
	All bool

	// WorkUnitKeys, if provided, is a list of specific work unit
	// keys to delete.  If this is given and All is false, then
	// these specific work units are deleted; if State is also
	// given, then each work unit must be in that state to be
	// deleted.
	WorkUnitKeys []string `mapstructure:"work_unit_keys"`

	// State, if provided, is one of the external Coordinate work
	// unit statuses, and all work units in this state are deleted.
	// If WorkUnitKeys is also provided then only those work units
	// will be deleted, and then only if in this state.
	State WorkUnitStatus
}

// DelWorkUnits deletes work units from an existing work spec.  If
// options is empty, this does nothing.  On success, returns the
// number of work units deleted.
func (jobs *JobServer) DelWorkUnits(workSpecName string, options map[string]interface{}) (int, string, error) {
	workSpec, err := jobs.Namespace.WorkSpec(workSpecName)
	var (
		count      int
		dwuOptions DelWorkUnitsOptions
		status     coordinate.WorkUnitStatus
	)
	if err == nil {
		err = decode(&dwuOptions, options)
	}
	if err == nil && !dwuOptions.All {
		status, err = translateWorkUnitStatus(dwuOptions.State)
	}
	if err == nil {
		var query coordinate.WorkUnitQuery
		if !dwuOptions.All {
			if dwuOptions.WorkUnitKeys != nil {
				query.Names = dwuOptions.WorkUnitKeys
			} else if status != coordinate.AnyStatus {
				query.Statuses = []coordinate.WorkUnitStatus{status}
			}
		}
		count, err = workSpec.DeleteWorkUnits(query)
	}
	return count, "", err
}

// Archive causes the system to clean up completed work units.  The
// system will keep up to a pre-specified limit of work units that
// have completed successfully, and will also remove work units that
// have completed successfully but are beyond a pre-specified age.
// The work units are deleted as in DelWorkUnits().  The return value
// is always nil.
//
// TODO(dmaze): Actually implement this.  This probably involves
// triggering a background task the system would need to do on its own
// in any case.  The observable effects of this are minimal, especially
// in a default/test configuration.
func (jobs *JobServer) Archive(options map[string]interface{}) (interface{}, error) {
	return nil, nil
}
<commit_msg>jobserver: Rewrite CountWorkUnits using new CountWorkUnitStatus<commit_after>package jobserver

import (
	"errors"
	"github.com/dmaze/goordinate/cborrpc"
	"github.com/dmaze/goordinate/coordinate"
	"github.com/mitchellh/mapstructure"
	"math"
	"reflect"
)

// AddWorkUnits adds any number of work units to a work spec.  Each oy
// the work units is a cborrpc.PythonTuple or slice containing a
// string with the work unit key, a dictionary with the work unit
// data, and an optional dictionary with additional metadata.
func (jobs *JobServer) AddWorkUnits(workSpecName string, workUnitKvp []interface{}) (bool, string, error) {
	spec, err := jobs.Namespace.WorkSpec(workSpecName)
	if err != nil {
		return false, "", err
	}

	// Unmarshal the work unit list into a []AddWorkUnitItem.
	// Fail now if any are invalid.
	items := make([]coordinate.AddWorkUnitItem, len(workUnitKvp))
	for i, kvp := range workUnitKvp {
		items[i], err = coordinate.ExtractAddWorkUnitItem(kvp)
		if err != nil {
			return false, "", err
		}
	}

	// Now go through and add them all
	for _, item := range items {
		_, err = spec.AddWorkUnit(item.Key, item.Data, item.Priority)
		if err != nil {
			// Again, Python coordinate expects to never see
			// a failure here?
			return false, "", err
		}
	}
	return true, "", nil
}

// GetWorkUnitsOptions contains unmarshaled options for GetWorkUnits().
type GetWorkUnitsOptions struct {
	// WorkUnitKeys contains a list of work unit keys to retrieve.
	// If this option is supplied, all other options are ignored.
	WorkUnitKeys []string `mapstructure:"work_unit_keys"`

	// State provides a list of states to query on.  If this is
	// provided then only work units in one of the specified states
	// will be returned.
	State []WorkUnitStatus

	// Start gives a starting point to iterate through the list of
	// work units.  It is the name of the last work unit returned
	// in the previous call to GetWorkUnits().  No work unit whose
	// name is lexicographically less than this will be returned.
	Start string

	// Limit specifies the maximum number of work units to return.
	// Defaults to 1000.
	Limit int
}

// gwuStateHook is a mapstructure decode hook that expands a single int
// or a PythonTuple into a slice of int (WorkUnitStatus).
func gwuStateHook(from reflect.Type, to reflect.Type, data interface{}) (interface{}, error) {
	// to must be []WorkUnitStatus
	if to.Kind() != reflect.Slice || to.Elem().Name() != "WorkUnitStatus" {
		return data, nil
	}
	switch value := data.(type) {
	case cborrpc.PythonTuple:
		// If from is a tuple, return its contents
		return value.Items, nil
	case WorkUnitStatus:
		// Package it into a slice
		return []WorkUnitStatus{value}, nil
	case int:
		// If from is an int, box it
		return []WorkUnitStatus{WorkUnitStatus(value)}, nil
	case uint64:
		return []WorkUnitStatus{WorkUnitStatus(value)}, nil
	default:
		// Otherwise, hope we can deal normally
		return data, nil
	}
}

// GetWorkUnits retrieves the keys and data dictionaries for some number
// of work units.  If options contains "work_unit_keys", those specific
// work units are retrieved; otherwise the work units are based on
// which of GetWorkUnitsOptions are present.
//
// On success, the return value is a slice of cborrpc.PythonTuple
// objects where each contains the work unit key as a byte slice and
// the data dictionary.
func (jobs *JobServer) GetWorkUnits(workSpecName string, options map[string]interface{}) ([]interface{}, string, error) {
	var workUnits map[string]coordinate.WorkUnit
	gwuOptions := GetWorkUnitsOptions{
		Limit: 1000,
	}

	spec, err := jobs.Namespace.WorkSpec(workSpecName)
	var decoder *mapstructure.Decoder
	if err == nil {
		config := mapstructure.DecoderConfig{
			DecodeHook: mapstructure.ComposeDecodeHookFunc(gwuStateHook, cborrpc.DecodeBytesAsString),
			Result:     &gwuOptions,
		}
		decoder, err = mapstructure.NewDecoder(&config)
	}
	if err == nil {
		err = decoder.Decode(options)
	}
	if err == nil {
		query := coordinate.WorkUnitQuery{
			Names: gwuOptions.WorkUnitKeys,
		}
		if gwuOptions.WorkUnitKeys == nil {
			query.PreviousName = gwuOptions.Start
			query.Limit = gwuOptions.Limit
		}
		if gwuOptions.WorkUnitKeys == nil && gwuOptions.State != nil {
			query.Statuses = make([]coordinate.WorkUnitStatus, len(gwuOptions.State))
			for i, state := range gwuOptions.State {
				query.Statuses[i], err = translateWorkUnitStatus(state)
				if err != nil {
					break
				}
			}
		}
		if err == nil {
			workUnits, err = spec.WorkUnits(query)
		}
	}
	if err != nil {
		return nil, "", err
	}
	// The marshalled result is a list of pairs of (key, data).
	var result []interface{}
	for name, unit := range workUnits {
		var data map[string]interface{}
		attempt, err := unit.ActiveAttempt()
		if err == nil && attempt != nil {
			data, err = attempt.Data()
		}
		if err == nil && data == nil {
			data, err = unit.Data()
		}
		if err != nil {
			return nil, "", err
		}
		tuple := cborrpc.PythonTuple{Items: []interface{}{[]byte(name), data}}
		result = append(result, tuple)
	}
	return result, "", nil
}

// GetWorkUnitStatus returns a summary status of zero or more work
// units in a single work spec.  On success, the returned list of
// dictionaries corresponds one-to-one with workUnitKeys.  If there is
// no such work unit, nil is in the list; otherwise each map contains
// keys "status", "expiration", "worker_id", and "traceback".
func (jobs *JobServer) GetWorkUnitStatus(workSpecName string, workUnitKeys []string) ([]map[string]interface{}, string, error) {
	spec, err := jobs.Namespace.WorkSpec(workSpecName)
	if err != nil {
		return nil, "", err
	}

	result := make([]map[string]interface{}, len(workUnitKeys))
	for i, key := range workUnitKeys {
		workUnit, err := spec.WorkUnit(key)
		if err != nil {
			return nil, "", err
		} else if workUnit == nil {
			result[i] = nil
		} else {
			r := make(map[string]interface{})
			status, attempt, err := workUnitStatus(workUnit)
			if err != nil {
				return nil, "", err
			}
			r["status"] = status
			if attempt != nil {
				r["worker_id"] = attempt.Worker().Name()
			}
			if status == Pending && attempt != nil {
				expiration, err := attempt.ExpirationTime()
				if err != nil {
					return nil, "", err
				}
				r["expiration"] = expiration.Unix()
			}
			if status == Failed && attempt != nil {
				data, err := attempt.Data()
				if err != nil {
					return nil, "", err
				}
				if traceback := data["traceback"]; traceback != nil {
					r["traceback"] = traceback
				}
			}
			result[i] = r
		}
	}
	return result, "", nil
}

// PrioritizeWorkUnitsOptions specifies which work units PrioritizeWorkUnits
// should adjust and how.
type PrioritizeWorkUnitsOptions struct {
	// WorkUnitKeys gives the names of the work units to reprioritize.
	// If not present, does nothing.
	WorkUnitKeys []string `mapstructure:"work_unit_keys"`

	// Priority sets an absolute priority.  If a NaN value, make a
	// change specified by Adjustment instead.
	Priority float64

	// Adjustment is added to the priorities of each of the work
	// units, if Priority is NaN.  If also a NaN value, do nothing.
	Adjustment float64
}

// PrioritizeWorkUnits changes the priorities of some number of work
// units.  The actual work units are in options["work_unit_keys"].  A
// higher priority results in the work units being scheduled sooner.
func (jobs *JobServer) PrioritizeWorkUnits(workSpecName string, options map[string]interface{}) (bool, string, error) {
	var (
		err      error
		query    coordinate.WorkUnitQuery
		workSpec coordinate.WorkSpec
	)
	pwuOptions := PrioritizeWorkUnitsOptions{
		Priority:   math.NaN(),
		Adjustment: math.NaN(),
	}
	workSpec, err = jobs.Namespace.WorkSpec(workSpecName)
	if err == nil {
		err = decode(&pwuOptions, options)
	}
	if err == nil && pwuOptions.WorkUnitKeys == nil {
		return false, "missing work_unit_keys", err
	}
	if err == nil {
		query.Names = pwuOptions.WorkUnitKeys
		if !math.IsNaN(pwuOptions.Priority) {
			err = workSpec.SetWorkUnitPriorities(query, pwuOptions.Priority)
		} else if !math.IsNaN(pwuOptions.Adjustment) {
			err = workSpec.AdjustWorkUnitPriorities(query, pwuOptions.Adjustment)
		}
	}
	return err == nil, "", err
}

// CountWorkUnits returns the number of work units in each status for
// a given work spec.
func (jobs *JobServer) CountWorkUnits(workSpecName string) (map[WorkUnitStatus]int, string, error) {
	workSpec, err := jobs.Namespace.WorkSpec(workSpecName)
	if err != nil {
		return nil, "", err
	}

	statuses, err := workSpec.CountWorkUnitStatus()
	if err != nil {
		return nil, "", err
	}

	result := make(map[WorkUnitStatus]int)
	for status, count := range statuses {
		var s WorkUnitStatus
		switch status {
		case coordinate.AvailableUnit:
			s = Available
		case coordinate.PendingUnit:
			s = Pending
		case coordinate.FinishedUnit:
			s = Finished
		case coordinate.FailedUnit:
			s = Failed
		}
		result[s] = count
	}
	return result, "", nil
}

// workUnitStatus extracts a summary of the status of a single work
// unit.  This produces its external coordinate status and the active
// attempt (if any) on success.
func workUnitStatus(workUnit coordinate.WorkUnit) (status WorkUnitStatus, attempt coordinate.Attempt, err error) {
	var attemptStatus coordinate.AttemptStatus
	attempt, err = workUnit.ActiveAttempt()
	if err == nil && attempt == nil {
		status = Available
		return
	}
	if err == nil {
		attemptStatus, err = attempt.Status()
	}
	if err == nil {
		switch attemptStatus {
		case coordinate.Pending:
			status = Pending
		case coordinate.Expired:
			status = Available
			attempt = nil
		case coordinate.Finished:
			status = Finished
		case coordinate.Failed:
			status = Failed
		case coordinate.Retryable:
			status = Available
			attempt = nil
		default:
			err = errors.New("unexpected attempt status")
		}
	}
	return
}

// DelWorkUnitsOptions specifies the options for DelWorkUnits.  The
// first of All, WorkUnitKeys, or State given defines the operation to
// perform.  If none of these are given, the zero value for this
// structure tells DelWorkUnits to do nothing.
type DelWorkUnitsOptions struct {
	// All, if set to true, directs DelWorkUnits to delete
	// all of the work units in its work spec.  If this is
	// provided, all other options are ignored.
	All bool

	// WorkUnitKeys, if provided, is a list of specific work unit
	// keys to delete.  If this is given and All is false, then
	// these specific work units are deleted; if State is also
	// given, then each work unit must be in that state to be
	// deleted.
	WorkUnitKeys []string `mapstructure:"work_unit_keys"`

	// State, if provided, is one of the external Coordinate work
	// unit statuses, and all work units in this state are deleted.
	// If WorkUnitKeys is also provided then only those work units
	// will be deleted, and then only if in this state.
	State WorkUnitStatus
}

// DelWorkUnits deletes work units from an existing work spec.  If
// options is empty, this does nothing.  On success, returns the
// number of work units deleted.
func (jobs *JobServer) DelWorkUnits(workSpecName string, options map[string]interface{}) (int, string, error) {
	workSpec, err := jobs.Namespace.WorkSpec(workSpecName)
	var (
		count      int
		dwuOptions DelWorkUnitsOptions
		status     coordinate.WorkUnitStatus
	)
	if err == nil {
		err = decode(&dwuOptions, options)
	}
	if err == nil && !dwuOptions.All {
		status, err = translateWorkUnitStatus(dwuOptions.State)
	}
	if err == nil {
		var query coordinate.WorkUnitQuery
		if !dwuOptions.All {
			if dwuOptions.WorkUnitKeys != nil {
				query.Names = dwuOptions.WorkUnitKeys
			} else if status != coordinate.AnyStatus {
				query.Statuses = []coordinate.WorkUnitStatus{status}
			}
		}
		count, err = workSpec.DeleteWorkUnits(query)
	}
	return count, "", err
}

// Archive causes the system to clean up completed work units.  The
// system will keep up to a pre-specified limit of work units that
// have completed successfully, and will also remove work units that
// have completed successfully but are beyond a pre-specified age.
// The work units are deleted as in DelWorkUnits().  The return value
// is always nil.
//
// TODO(dmaze): Actually implement this.  This probably involves
// triggering a background task the system would need to do on its own
// in any case.  The observable effects of this are minimal, especially
// in a default/test configuration.
func (jobs *JobServer) Archive(options map[string]interface{}) (interface{}, error) {
	return nil, nil
}
<|endoftext|>


<commit_before>// Copyright 2016 CoreOS, Inc.
//
// 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 listener

import (
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"os"
	"time"
)

type listenerUnix struct {
	net.Listener
}

func (lu *listenerUnix) Close() error {
	if err := os.RemoveAll(lu.Addr().String()); err != nil {
		return err
	}
	return lu.Listener.Close()
}

// NewListenerUnix returns new net.Listener with unix socket
// (unix sockets via unix://address).
//
// (etcd pkg.transport.NewUnixListener)
func NewListenerUnix(addr string) (net.Listener, error) {
	if err := os.RemoveAll(addr); err != nil {
		return nil, err
	}
	l, err := net.Listen("unix", addr)
	if err != nil {
		return nil, err
	}
	return &listenerUnix{l}, nil
}

// NewListener returns new net.Listener based on the scheme and tls.Config.
//
// (etcd pkg.transport.NewListener)
func NewListener(scheme, addr string, tlsConfig *tls.Config) (l net.Listener, err error) {
	switch scheme {
	case "unix", "unixs":
		l, err = NewListenerUnix(addr)
		if err != nil {
			return
		}

	case "http", "https":
		l, err = net.Listen("tcp", addr)
		if err != nil {
			return
		}

	default:
		return nil, fmt.Errorf("%q is not supported", scheme)
	}

	if scheme != "https" && scheme != "unixs" { // no need TLS
		return
	}
	if tlsConfig == nil { // need TLS, but empty config
		return nil, fmt.Errorf("cannot listen on TLS for %s: KeyFile and CertFile are not presented", scheme+"://"+addr)
	}
	return tls.NewListener(l, tlsConfig), nil
}

type listenerStoppable struct {
	net.Listener
	stopc <-chan struct{}
}

// NewListenerStoppable returns stoppable net.Listener.
func NewListenerStoppable(scheme, addr string, tlsConfig *tls.Config, stopc <-chan struct{}) (net.Listener, error) {
	ln, err := NewListener(scheme, addr, tlsConfig)
	if err != nil {
		return nil, err
	}
	ls := &listenerStoppable{
		Listener: ln,
		stopc:    stopc,
	}
	return ls, nil
}

// ErrListenerStopped is returned when the listener is stopped.
var ErrListenerStopped = errors.New("listener stopped")

func (ln *listenerStoppable) Accept() (net.Conn, error) {
	connc, errc := make(chan net.Conn, 1), make(chan error)
	go func() {
		conn, err := ln.Listener.Accept() // (X) ln.Accept()
		if err != nil {
			errc <- err
			return
		}
		connc <- conn
	}()

	select {
	case <-ln.stopc:
		return nil, ErrListenerStopped
	case err := <-errc:
		return nil, err
	case conn := <-connc:
		tc, ok := conn.(*net.TCPConn)
		if ok {
			tc.SetKeepAlive(true)
			tc.SetKeepAlivePeriod(3 * time.Minute)
		}
		return conn, nil
	}
}
<commit_msg>listener: close listener<commit_after>// Copyright 2016 CoreOS, Inc.
//
// 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 listener

import (
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"os"
	"time"
)

type listenerUnix struct {
	net.Listener
}

func (lu *listenerUnix) Close() error {
	if err := os.RemoveAll(lu.Addr().String()); err != nil {
		return err
	}
	return lu.Listener.Close()
}

// NewListenerUnix returns new net.Listener with unix socket
// (unix sockets via unix://address).
//
// (etcd pkg.transport.NewUnixListener)
func NewListenerUnix(addr string) (net.Listener, error) {
	if err := os.RemoveAll(addr); err != nil {
		return nil, err
	}
	l, err := net.Listen("unix", addr)
	if err != nil {
		return nil, err
	}
	return &listenerUnix{l}, nil
}

// NewListener returns new net.Listener based on the scheme and tls.Config.
//
// (etcd pkg.transport.NewListener)
func NewListener(scheme, addr string, tlsConfig *tls.Config) (l net.Listener, err error) {
	switch scheme {
	case "unix", "unixs":
		l, err = NewListenerUnix(addr)
		if err != nil {
			return
		}

	case "http", "https":
		l, err = net.Listen("tcp", addr)
		if err != nil {
			return
		}

	default:
		return nil, fmt.Errorf("%q is not supported", scheme)
	}

	if scheme != "https" && scheme != "unixs" { // no need TLS
		return
	}
	if tlsConfig == nil { // need TLS, but empty config
		return nil, fmt.Errorf("cannot listen on TLS for %s: KeyFile and CertFile are not presented", scheme+"://"+addr)
	}
	return tls.NewListener(l, tlsConfig), nil
}

type listenerStoppable struct {
	net.Listener
	stopc <-chan struct{}
}

// NewListenerStoppable returns stoppable net.Listener.
func NewListenerStoppable(scheme, addr string, tlsConfig *tls.Config, stopc <-chan struct{}) (net.Listener, error) {
	ln, err := NewListener(scheme, addr, tlsConfig)
	if err != nil {
		return nil, err
	}
	ls := &listenerStoppable{
		Listener: ln,
		stopc:    stopc,
	}
	return ls, nil
}

// ErrListenerStopped is returned when the listener is stopped.
var ErrListenerStopped = errors.New("listener stopped")

func (ln *listenerStoppable) Accept() (net.Conn, error) {
	connc, errc := make(chan net.Conn, 1), make(chan error)
	go func() {
		conn, err := ln.Listener.Accept() // (X) ln.Accept()
		if err != nil {
			errc <- err
			return
		}
		connc <- conn
	}()

	select {
	case <-ln.stopc:
		if cerr := ln.Close(); cerr != nil {
			return nil, cerr
		}
		return nil, ErrListenerStopped
	case err := <-errc:
		return nil, err
	case conn := <-connc:
		tc, ok := conn.(*net.TCPConn)
		if ok {
			tc.SetKeepAlive(true)
			tc.SetKeepAlivePeriod(3 * time.Minute)
		}
		return conn, nil
	}
}
<|endoftext|>