<commit_before>package backends

import (
	"bytes"
	"context"
	"errors"
	"net"
	"strconv"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestSend(t *testing.T) {
	dc := dummyConn{}
	sender := Sender{
		ConnFactory: func() (net.Conn, error) {
			return &dc, nil
		},
		Sink: make(chan Stream),
		BufPool: sync.Pool{
			New: func() interface{} {
				return new(bytes.Buffer)
			},
		},
	}
	var wg sync.WaitGroup
	defer wg.Wait()
	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(10*time.Second))
	defer cancel()
	wg.Add(1)
	go func() {
		defer wg.Done()
		if err := sender.Run(ctx); err != nil && err != context.Canceled && err != context.DeadlineExceeded {
			t.Error(err)
		}
	}()
	var wgTest sync.WaitGroup
	for i := 0; i <= 4; i++ {
		wgTest.Add(1)
		t.Run(strconv.Itoa(i), func(t *testing.T) {
			sink := make(chan *bytes.Buffer, i)
			sender.Sink <- Stream{
				Cb: func(errs []error) {
					defer wgTest.Done()
					for _, e := range errs {
						assert.NoError(t, e)
					}
				},
				Buf: sink,
			}
			for x := 0; x < i; x++ {
				sink <- bytes.NewBuffer([]byte{byte(x)})
			}
			close(sink)
		})
	}
	wgTest.Wait()
	assert.Equal(t, []byte{0x0, 0x0, 0x1, 0x0, 0x1, 0x2, 0x0, 0x1, 0x2, 0x3}, dc.buf.Bytes())
}

type dummyConn struct {
	buf      bytes.Buffer
	isClosed bool
}

func (c *dummyConn) Read(b []byte) (int, error) {
	return 0, errors.New("asdasd")
}

func (c *dummyConn) Write(b []byte) (int, error) {
	if c.isClosed {
		panic("closed")
	}
	return c.buf.Write(b)
}

func (c *dummyConn) Close() error {
	c.isClosed = true
	return nil
}

func (c *dummyConn) LocalAddr() net.Addr {
	return nil
}

func (c *dummyConn) RemoteAddr() net.Addr {
	return nil
}

func (c *dummyConn) SetDeadline(t time.Time) error {
	return nil
}

func (c *dummyConn) SetReadDeadline(t time.Time) error {
	return nil
}

func (c *dummyConn) SetWriteDeadline(t time.Time) error {
	return nil
}
<commit_msg>Simplify<commit_after>package backends

import (
	"bytes"
	"context"
	"errors"
	"net"
	"strconv"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func TestSend(t *testing.T) {
	dc := dummyConn{}
	sender := Sender{
		ConnFactory: func() (net.Conn, error) {
			return &dc, nil
		},
		Sink: make(chan Stream),
		BufPool: sync.Pool{
			New: func() interface{} {
				return new(bytes.Buffer)
			},
		},
	}
	var wg sync.WaitGroup
	defer wg.Wait()
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	wg.Add(1)
	go func() {
		defer wg.Done()
		if err := sender.Run(ctx); err != nil && err != context.Canceled && err != context.DeadlineExceeded {
			t.Error(err)
		}
	}()
	var wgTest sync.WaitGroup
	for i := 0; i <= 4; i++ {
		wgTest.Add(1)
		t.Run(strconv.Itoa(i), func(t *testing.T) {
			sink := make(chan *bytes.Buffer, i)
			sender.Sink <- Stream{
				Cb: func(errs []error) {
					defer wgTest.Done()
					for _, e := range errs {
						assert.NoError(t, e)
					}
				},
				Buf: sink,
			}
			for x := 0; x < i; x++ {
				sink <- bytes.NewBuffer([]byte{byte(x)})
			}
			close(sink)
		})
	}
	wgTest.Wait()
	assert.Equal(t, []byte{0x0, 0x0, 0x1, 0x0, 0x1, 0x2, 0x0, 0x1, 0x2, 0x3}, dc.buf.Bytes())
}

type dummyConn struct {
	buf      bytes.Buffer
	isClosed bool
}

func (c *dummyConn) Read(b []byte) (int, error) {
	return 0, errors.New("asdasd")
}

func (c *dummyConn) Write(b []byte) (int, error) {
	if c.isClosed {
		panic("closed")
	}
	return c.buf.Write(b)
}

func (c *dummyConn) Close() error {
	c.isClosed = true
	return nil
}

func (c *dummyConn) LocalAddr() net.Addr {
	return nil
}

func (c *dummyConn) RemoteAddr() net.Addr {
	return nil
}

func (c *dummyConn) SetDeadline(t time.Time) error {
	return nil
}

func (c *dummyConn) SetReadDeadline(t time.Time) error {
	return nil
}

func (c *dummyConn) SetWriteDeadline(t time.Time) error {
	return nil
}
<|endoftext|>


<commit_before>package hcs

import (
	"encoding/json"
	"io"
	"sync"
	"syscall"
	"time"

	"github.com/Microsoft/hcsshim/internal/interop"
	"github.com/Microsoft/hcsshim/internal/logfields"
	"github.com/sirupsen/logrus"
)

// ContainerError is an error encountered in HCS
type Process struct {
	handleLock     sync.RWMutex
	handle         hcsProcess
	processID      int
	system         *System
	cachedPipes    *cachedPipes
	callbackNumber uintptr

	logctx logrus.Fields

	closedWaitOnce sync.Once
	waitBlock      chan struct{}
	waitError      error
}

func newProcess(process hcsProcess, processID int, computeSystem *System) *Process {
	return &Process{
		handle:    process,
		processID: processID,
		system:    computeSystem,
		logctx: logrus.Fields{
			logfields.ContainerID: computeSystem.ID(),
			logfields.ProcessID:   processID,
		},
		waitBlock: make(chan struct{}),
	}
}

type cachedPipes struct {
	stdIn  syscall.Handle
	stdOut syscall.Handle
	stdErr syscall.Handle
}

type processModifyRequest struct {
	Operation   string
	ConsoleSize *consoleSize `json:",omitempty"`
	CloseHandle *closeHandle `json:",omitempty"`
}

type consoleSize struct {
	Height uint16
	Width  uint16
}

type closeHandle struct {
	Handle string
}

type ProcessStatus struct {
	ProcessID      uint32
	Exited         bool
	ExitCode       uint32
	LastWaitResult int32
}

const (
	stdIn  string = "StdIn"
	stdOut string = "StdOut"
	stdErr string = "StdErr"
)

const (
	modifyConsoleSize string = "ConsoleSize"
	modifyCloseHandle string = "CloseHandle"
)

// Pid returns the process ID of the process within the container.
func (process *Process) Pid() int {
	return process.processID
}

// SystemID returns the ID of the process's compute system.
func (process *Process) SystemID() string {
	return process.system.ID()
}

func (process *Process) logOperationBegin(operation string) {
	logOperationBegin(
		process.logctx,
		operation+" - Begin Operation")
}

func (process *Process) logOperationEnd(operation string, err error) {
	var result string
	if err == nil {
		result = "Success"
	} else {
		result = "Error"
	}

	logOperationEnd(
		process.logctx,
		operation+" - End Operation - "+result,
		err)
}

// Signal signals the process with `options`.
//
// For LCOW `guestrequest.SignalProcessOptionsLCOW`.
//
// For WCOW `guestrequest.SignalProcessOptionsWCOW`.
func (process *Process) Signal(options interface{}) (err error) {
	process.handleLock.RLock()
	defer process.handleLock.RUnlock()

	operation := "hcsshim::Process::Signal"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	if process.handle == 0 {
		return makeProcessError(process, operation, ErrAlreadyClosed, nil)
	}

	optionsb, err := json.Marshal(options)
	if err != nil {
		return err
	}

	optionsStr := string(optionsb)

	var resultp *uint16
	syscallWatcher(process.logctx, func() {
		err = hcsSignalProcess(process.handle, optionsStr, &resultp)
	})
	events := processHcsResult(resultp)
	if err != nil {
		return makeProcessError(process, operation, err, events)
	}

	return nil
}

// Kill signals the process to terminate but does not wait for it to finish terminating.
func (process *Process) Kill() (err error) {
	process.handleLock.RLock()
	defer process.handleLock.RUnlock()

	operation := "hcsshim::Process::Kill"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	if process.handle == 0 {
		return makeProcessError(process, operation, ErrAlreadyClosed, nil)
	}

	var resultp *uint16
	syscallWatcher(process.logctx, func() {
		err = hcsTerminateProcess(process.handle, &resultp)
	})
	events := processHcsResult(resultp)
	if err != nil {
		return makeProcessError(process, operation, err, events)
	}

	return nil
}

// waitBackground waits for the process exit notification. Once received sets
// `process.waitError` (if any) and unblocks all `Wait` and `WaitTimeout` calls.
//
// This MUST be called exactly once per `process.handle` but `Wait` and
// `WaitTimeout` are safe to call multiple times.
func (process *Process) waitBackground() {
	process.waitError = waitForNotification(process.callbackNumber, hcsNotificationProcessExited, nil)
	process.closedWaitOnce.Do(func() {
		close(process.waitBlock)
	})
}

// Wait waits for the process to exit. If the process has already exited returns
// the pervious error (if any).
func (process *Process) Wait() (err error) {
	operation := "hcsshim::Process::Wait"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	<-process.waitBlock
	if process.waitError != nil {
		return makeProcessError(process, operation, err, nil)
	}
	return nil
}

// WaitTimeout waits for the process to exit or the duration to elapse. If the
// process has already exited returns the pervious error (if any). If a timeout
// occurs returns `ErrTimeout`.
func (process *Process) WaitTimeout(timeout time.Duration) (err error) {
	operation := "hcssshim::Process::WaitTimeout"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	select {
	case <-process.waitBlock:
		if process.waitError != nil {
			return makeProcessError(process, operation, process.waitError, nil)
		}
		return nil
	case <-time.After(timeout):
		return makeProcessError(process, operation, ErrTimeout, nil)
	}
}

// ResizeConsole resizes the console of the process.
func (process *Process) ResizeConsole(width, height uint16) (err error) {
	process.handleLock.RLock()
	defer process.handleLock.RUnlock()

	operation := "hcsshim::Process::ResizeConsole"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	if process.handle == 0 {
		return makeProcessError(process, operation, ErrAlreadyClosed, nil)
	}

	modifyRequest := processModifyRequest{
		Operation: modifyConsoleSize,
		ConsoleSize: &consoleSize{
			Height: height,
			Width:  width,
		},
	}

	modifyRequestb, err := json.Marshal(modifyRequest)
	if err != nil {
		return err
	}

	modifyRequestStr := string(modifyRequestb)

	var resultp *uint16
	err = hcsModifyProcess(process.handle, modifyRequestStr, &resultp)
	events := processHcsResult(resultp)
	if err != nil {
		return makeProcessError(process, operation, err, events)
	}

	return nil
}

func (process *Process) Properties() (_ *ProcessStatus, err error) {
	process.handleLock.RLock()
	defer process.handleLock.RUnlock()

	operation := "hcsshim::Process::Properties"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	if process.handle == 0 {
		return nil, makeProcessError(process, operation, ErrAlreadyClosed, nil)
	}

	var (
		resultp     *uint16
		propertiesp *uint16
	)
	syscallWatcher(process.logctx, func() {
		err = hcsGetProcessProperties(process.handle, &propertiesp, &resultp)
	})
	events := processHcsResult(resultp)
	if err != nil {
		return nil, makeProcessError(process, operation, err, events)
	}

	if propertiesp == nil {
		return nil, ErrUnexpectedValue
	}
	propertiesRaw := interop.ConvertAndFreeCoTaskMemBytes(propertiesp)

	properties := &ProcessStatus{}
	if err := json.Unmarshal(propertiesRaw, properties); err != nil {
		return nil, makeProcessError(process, operation, err, nil)
	}

	return properties, nil
}

// ExitCode returns the exit code of the process. The process must have
// already terminated.
func (process *Process) ExitCode() (_ int, err error) {
	operation := "hcsshim::Process::ExitCode"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	properties, err := process.Properties()
	if err != nil {
		return -1, makeProcessError(process, operation, err, nil)
	}

	if properties.Exited == false {
		return -1, makeProcessError(process, operation, ErrInvalidProcessState, nil)
	}

	if properties.LastWaitResult != 0 {
		logrus.WithFields(logrus.Fields{
			logfields.ContainerID: process.SystemID(),
			logfields.ProcessID:   process.processID,
			"wait-result":         properties.LastWaitResult,
		}).Warn("hcsshim::Process::ExitCode - Non-zero last wait result")
		return -1, nil
	}

	return int(properties.ExitCode), nil
}

// Stdio returns the stdin, stdout, and stderr pipes, respectively. Closing
// these pipes does not close the underlying pipes; it should be possible to
// call this multiple times to get multiple interfaces.
func (process *Process) Stdio() (_ io.WriteCloser, _ io.ReadCloser, _ io.ReadCloser, err error) {
	process.handleLock.RLock()
	defer process.handleLock.RUnlock()

	operation := "hcsshim::Process::Stdio"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	if process.handle == 0 {
		return nil, nil, nil, makeProcessError(process, operation, ErrAlreadyClosed, nil)
	}

	var stdIn, stdOut, stdErr syscall.Handle

	if process.cachedPipes == nil {
		var (
			processInfo hcsProcessInformation
			resultp     *uint16
		)
		err = hcsGetProcessInfo(process.handle, &processInfo, &resultp)
		events := processHcsResult(resultp)
		if err != nil {
			return nil, nil, nil, makeProcessError(process, operation, err, events)
		}

		stdIn, stdOut, stdErr = processInfo.StdInput, processInfo.StdOutput, processInfo.StdError
	} else {
		// Use cached pipes
		stdIn, stdOut, stdErr = process.cachedPipes.stdIn, process.cachedPipes.stdOut, process.cachedPipes.stdErr

		// Invalidate the cache
		process.cachedPipes = nil
	}

	pipes, err := makeOpenFiles([]syscall.Handle{stdIn, stdOut, stdErr})
	if err != nil {
		return nil, nil, nil, makeProcessError(process, operation, err, nil)
	}

	return pipes[0], pipes[1], pipes[2], nil
}

// CloseStdin closes the write side of the stdin pipe so that the process is
// notified on the read side that there is no more data in stdin.
func (process *Process) CloseStdin() (err error) {
	process.handleLock.RLock()
	defer process.handleLock.RUnlock()

	operation := "hcsshim::Process::CloseStdin"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	if process.handle == 0 {
		return makeProcessError(process, operation, ErrAlreadyClosed, nil)
	}

	modifyRequest := processModifyRequest{
		Operation: modifyCloseHandle,
		CloseHandle: &closeHandle{
			Handle: stdIn,
		},
	}

	modifyRequestb, err := json.Marshal(modifyRequest)
	if err != nil {
		return err
	}

	modifyRequestStr := string(modifyRequestb)

	var resultp *uint16
	err = hcsModifyProcess(process.handle, modifyRequestStr, &resultp)
	events := processHcsResult(resultp)
	if err != nil {
		return makeProcessError(process, operation, err, events)
	}

	return nil
}

// Close cleans up any state associated with the process but does not kill
// or wait on it.
func (process *Process) Close() (err error) {
	process.handleLock.Lock()
	defer process.handleLock.Unlock()

	operation := "hcsshim::Process::Close"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	// Don't double free this
	if process.handle == 0 {
		return nil
	}

	if err = process.unregisterCallback(); err != nil {
		return makeProcessError(process, operation, err, nil)
	}

	if err = hcsCloseProcess(process.handle); err != nil {
		return makeProcessError(process, operation, err, nil)
	}

	process.handle = 0
	process.closedWaitOnce.Do(func() {
		close(process.waitBlock)
	})

	return nil
}

func (process *Process) registerCallback() error {
	context := &notifcationWatcherContext{
		channels:  newProcessChannels(),
		systemID:  process.SystemID(),
		processID: process.processID,
	}

	callbackMapLock.Lock()
	callbackNumber := nextCallback
	nextCallback++
	callbackMap[callbackNumber] = context
	callbackMapLock.Unlock()

	var callbackHandle hcsCallback
	err := hcsRegisterProcessCallback(process.handle, notificationWatcherCallback, callbackNumber, &callbackHandle)
	if err != nil {
		return err
	}
	context.handle = callbackHandle
	process.callbackNumber = callbackNumber

	return nil
}

func (process *Process) unregisterCallback() error {
	callbackNumber := process.callbackNumber

	callbackMapLock.RLock()
	context := callbackMap[callbackNumber]
	callbackMapLock.RUnlock()

	if context == nil {
		return nil
	}

	handle := context.handle

	if handle == 0 {
		return nil
	}

	// hcsUnregisterProcessCallback has its own syncronization
	// to wait for all callbacks to complete. We must NOT hold the callbackMapLock.
	err := hcsUnregisterProcessCallback(handle)
	if err != nil {
		return err
	}

	closeChannels(context.channels)

	callbackMapLock.Lock()
	delete(callbackMap, callbackNumber)
	callbackMapLock.Unlock()

	handle = 0

	return nil
}
<commit_msg>Fix panic on hcs/process.Close() with nil error<commit_after>package hcs

import (
	"encoding/json"
	"io"
	"sync"
	"syscall"
	"time"

	"github.com/Microsoft/hcsshim/internal/interop"
	"github.com/Microsoft/hcsshim/internal/logfields"
	"github.com/sirupsen/logrus"
)

// ContainerError is an error encountered in HCS
type Process struct {
	handleLock     sync.RWMutex
	handle         hcsProcess
	processID      int
	system         *System
	cachedPipes    *cachedPipes
	callbackNumber uintptr

	logctx logrus.Fields

	closedWaitOnce sync.Once
	waitBlock      chan struct{}
	waitError      error
}

func newProcess(process hcsProcess, processID int, computeSystem *System) *Process {
	return &Process{
		handle:    process,
		processID: processID,
		system:    computeSystem,
		logctx: logrus.Fields{
			logfields.ContainerID: computeSystem.ID(),
			logfields.ProcessID:   processID,
		},
		waitBlock: make(chan struct{}),
	}
}

type cachedPipes struct {
	stdIn  syscall.Handle
	stdOut syscall.Handle
	stdErr syscall.Handle
}

type processModifyRequest struct {
	Operation   string
	ConsoleSize *consoleSize `json:",omitempty"`
	CloseHandle *closeHandle `json:",omitempty"`
}

type consoleSize struct {
	Height uint16
	Width  uint16
}

type closeHandle struct {
	Handle string
}

type ProcessStatus struct {
	ProcessID      uint32
	Exited         bool
	ExitCode       uint32
	LastWaitResult int32
}

const (
	stdIn  string = "StdIn"
	stdOut string = "StdOut"
	stdErr string = "StdErr"
)

const (
	modifyConsoleSize string = "ConsoleSize"
	modifyCloseHandle string = "CloseHandle"
)

// Pid returns the process ID of the process within the container.
func (process *Process) Pid() int {
	return process.processID
}

// SystemID returns the ID of the process's compute system.
func (process *Process) SystemID() string {
	return process.system.ID()
}

func (process *Process) logOperationBegin(operation string) {
	logOperationBegin(
		process.logctx,
		operation+" - Begin Operation")
}

func (process *Process) logOperationEnd(operation string, err error) {
	var result string
	if err == nil {
		result = "Success"
	} else {
		result = "Error"
	}

	logOperationEnd(
		process.logctx,
		operation+" - End Operation - "+result,
		err)
}

// Signal signals the process with `options`.
//
// For LCOW `guestrequest.SignalProcessOptionsLCOW`.
//
// For WCOW `guestrequest.SignalProcessOptionsWCOW`.
func (process *Process) Signal(options interface{}) (err error) {
	process.handleLock.RLock()
	defer process.handleLock.RUnlock()

	operation := "hcsshim::Process::Signal"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	if process.handle == 0 {
		return makeProcessError(process, operation, ErrAlreadyClosed, nil)
	}

	optionsb, err := json.Marshal(options)
	if err != nil {
		return err
	}

	optionsStr := string(optionsb)

	var resultp *uint16
	syscallWatcher(process.logctx, func() {
		err = hcsSignalProcess(process.handle, optionsStr, &resultp)
	})
	events := processHcsResult(resultp)
	if err != nil {
		return makeProcessError(process, operation, err, events)
	}

	return nil
}

// Kill signals the process to terminate but does not wait for it to finish terminating.
func (process *Process) Kill() (err error) {
	process.handleLock.RLock()
	defer process.handleLock.RUnlock()

	operation := "hcsshim::Process::Kill"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	if process.handle == 0 {
		return makeProcessError(process, operation, ErrAlreadyClosed, nil)
	}

	var resultp *uint16
	syscallWatcher(process.logctx, func() {
		err = hcsTerminateProcess(process.handle, &resultp)
	})
	events := processHcsResult(resultp)
	if err != nil {
		return makeProcessError(process, operation, err, events)
	}

	return nil
}

// waitBackground waits for the process exit notification. Once received sets
// `process.waitError` (if any) and unblocks all `Wait` and `WaitTimeout` calls.
//
// This MUST be called exactly once per `process.handle` but `Wait` and
// `WaitTimeout` are safe to call multiple times.
func (process *Process) waitBackground() {
	process.waitError = waitForNotification(process.callbackNumber, hcsNotificationProcessExited, nil)
	process.closedWaitOnce.Do(func() {
		close(process.waitBlock)
	})
}

// Wait waits for the process to exit. If the process has already exited returns
// the pervious error (if any).
func (process *Process) Wait() (err error) {
	operation := "hcsshim::Process::Wait"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	<-process.waitBlock
	if process.waitError != nil {
		return makeProcessError(process, operation, process.waitError, nil)
	}
	return nil
}

// WaitTimeout waits for the process to exit or the duration to elapse. If the
// process has already exited returns the pervious error (if any). If a timeout
// occurs returns `ErrTimeout`.
func (process *Process) WaitTimeout(timeout time.Duration) (err error) {
	operation := "hcssshim::Process::WaitTimeout"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	select {
	case <-process.waitBlock:
		if process.waitError != nil {
			return makeProcessError(process, operation, process.waitError, nil)
		}
		return nil
	case <-time.After(timeout):
		return makeProcessError(process, operation, ErrTimeout, nil)
	}
}

// ResizeConsole resizes the console of the process.
func (process *Process) ResizeConsole(width, height uint16) (err error) {
	process.handleLock.RLock()
	defer process.handleLock.RUnlock()

	operation := "hcsshim::Process::ResizeConsole"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	if process.handle == 0 {
		return makeProcessError(process, operation, ErrAlreadyClosed, nil)
	}

	modifyRequest := processModifyRequest{
		Operation: modifyConsoleSize,
		ConsoleSize: &consoleSize{
			Height: height,
			Width:  width,
		},
	}

	modifyRequestb, err := json.Marshal(modifyRequest)
	if err != nil {
		return err
	}

	modifyRequestStr := string(modifyRequestb)

	var resultp *uint16
	err = hcsModifyProcess(process.handle, modifyRequestStr, &resultp)
	events := processHcsResult(resultp)
	if err != nil {
		return makeProcessError(process, operation, err, events)
	}

	return nil
}

func (process *Process) Properties() (_ *ProcessStatus, err error) {
	process.handleLock.RLock()
	defer process.handleLock.RUnlock()

	operation := "hcsshim::Process::Properties"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	if process.handle == 0 {
		return nil, makeProcessError(process, operation, ErrAlreadyClosed, nil)
	}

	var (
		resultp     *uint16
		propertiesp *uint16
	)
	syscallWatcher(process.logctx, func() {
		err = hcsGetProcessProperties(process.handle, &propertiesp, &resultp)
	})
	events := processHcsResult(resultp)
	if err != nil {
		return nil, makeProcessError(process, operation, err, events)
	}

	if propertiesp == nil {
		return nil, ErrUnexpectedValue
	}
	propertiesRaw := interop.ConvertAndFreeCoTaskMemBytes(propertiesp)

	properties := &ProcessStatus{}
	if err := json.Unmarshal(propertiesRaw, properties); err != nil {
		return nil, makeProcessError(process, operation, err, nil)
	}

	return properties, nil
}

// ExitCode returns the exit code of the process. The process must have
// already terminated.
func (process *Process) ExitCode() (_ int, err error) {
	operation := "hcsshim::Process::ExitCode"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	properties, err := process.Properties()
	if err != nil {
		return -1, makeProcessError(process, operation, err, nil)
	}

	if properties.Exited == false {
		return -1, makeProcessError(process, operation, ErrInvalidProcessState, nil)
	}

	if properties.LastWaitResult != 0 {
		logrus.WithFields(logrus.Fields{
			logfields.ContainerID: process.SystemID(),
			logfields.ProcessID:   process.processID,
			"wait-result":         properties.LastWaitResult,
		}).Warn("hcsshim::Process::ExitCode - Non-zero last wait result")
		return -1, nil
	}

	return int(properties.ExitCode), nil
}

// Stdio returns the stdin, stdout, and stderr pipes, respectively. Closing
// these pipes does not close the underlying pipes; it should be possible to
// call this multiple times to get multiple interfaces.
func (process *Process) Stdio() (_ io.WriteCloser, _ io.ReadCloser, _ io.ReadCloser, err error) {
	process.handleLock.RLock()
	defer process.handleLock.RUnlock()

	operation := "hcsshim::Process::Stdio"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	if process.handle == 0 {
		return nil, nil, nil, makeProcessError(process, operation, ErrAlreadyClosed, nil)
	}

	var stdIn, stdOut, stdErr syscall.Handle

	if process.cachedPipes == nil {
		var (
			processInfo hcsProcessInformation
			resultp     *uint16
		)
		err = hcsGetProcessInfo(process.handle, &processInfo, &resultp)
		events := processHcsResult(resultp)
		if err != nil {
			return nil, nil, nil, makeProcessError(process, operation, err, events)
		}

		stdIn, stdOut, stdErr = processInfo.StdInput, processInfo.StdOutput, processInfo.StdError
	} else {
		// Use cached pipes
		stdIn, stdOut, stdErr = process.cachedPipes.stdIn, process.cachedPipes.stdOut, process.cachedPipes.stdErr

		// Invalidate the cache
		process.cachedPipes = nil
	}

	pipes, err := makeOpenFiles([]syscall.Handle{stdIn, stdOut, stdErr})
	if err != nil {
		return nil, nil, nil, makeProcessError(process, operation, err, nil)
	}

	return pipes[0], pipes[1], pipes[2], nil
}

// CloseStdin closes the write side of the stdin pipe so that the process is
// notified on the read side that there is no more data in stdin.
func (process *Process) CloseStdin() (err error) {
	process.handleLock.RLock()
	defer process.handleLock.RUnlock()

	operation := "hcsshim::Process::CloseStdin"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	if process.handle == 0 {
		return makeProcessError(process, operation, ErrAlreadyClosed, nil)
	}

	modifyRequest := processModifyRequest{
		Operation: modifyCloseHandle,
		CloseHandle: &closeHandle{
			Handle: stdIn,
		},
	}

	modifyRequestb, err := json.Marshal(modifyRequest)
	if err != nil {
		return err
	}

	modifyRequestStr := string(modifyRequestb)

	var resultp *uint16
	err = hcsModifyProcess(process.handle, modifyRequestStr, &resultp)
	events := processHcsResult(resultp)
	if err != nil {
		return makeProcessError(process, operation, err, events)
	}

	return nil
}

// Close cleans up any state associated with the process but does not kill
// or wait on it.
func (process *Process) Close() (err error) {
	process.handleLock.Lock()
	defer process.handleLock.Unlock()

	operation := "hcsshim::Process::Close"
	process.logOperationBegin(operation)
	defer func() { process.logOperationEnd(operation, err) }()

	// Don't double free this
	if process.handle == 0 {
		return nil
	}

	if err = process.unregisterCallback(); err != nil {
		return makeProcessError(process, operation, err, nil)
	}

	if err = hcsCloseProcess(process.handle); err != nil {
		return makeProcessError(process, operation, err, nil)
	}

	process.handle = 0
	process.closedWaitOnce.Do(func() {
		close(process.waitBlock)
	})

	return nil
}

func (process *Process) registerCallback() error {
	context := &notifcationWatcherContext{
		channels:  newProcessChannels(),
		systemID:  process.SystemID(),
		processID: process.processID,
	}

	callbackMapLock.Lock()
	callbackNumber := nextCallback
	nextCallback++
	callbackMap[callbackNumber] = context
	callbackMapLock.Unlock()

	var callbackHandle hcsCallback
	err := hcsRegisterProcessCallback(process.handle, notificationWatcherCallback, callbackNumber, &callbackHandle)
	if err != nil {
		return err
	}
	context.handle = callbackHandle
	process.callbackNumber = callbackNumber

	return nil
}

func (process *Process) unregisterCallback() error {
	callbackNumber := process.callbackNumber

	callbackMapLock.RLock()
	context := callbackMap[callbackNumber]
	callbackMapLock.RUnlock()

	if context == nil {
		return nil
	}

	handle := context.handle

	if handle == 0 {
		return nil
	}

	// hcsUnregisterProcessCallback has its own syncronization
	// to wait for all callbacks to complete. We must NOT hold the callbackMapLock.
	err := hcsUnregisterProcessCallback(handle)
	if err != nil {
		return err
	}

	closeChannels(context.channels)

	callbackMapLock.Lock()
	delete(callbackMap, callbackNumber)
	callbackMapLock.Unlock()

	handle = 0

	return nil
}
<|endoftext|>


<commit_before>package internal

import (
	"encoding/xml"
	"io"
)

// Ref: tealeg/xlsx/xmlWorkbook.go
// I just keep the part I want.

func readWorkbookXML(rc io.ReadCloser) (*xlsxWorkbook, error) {
	var err error
	workbook := new(xlsxWorkbook)
	decoder := xml.NewDecoder(rc)
	err = decoder.Decode(workbook)
	if err != nil {
		return nil, err
	}
	return workbook, nil
}

// xlsxWorkbook directly maps the workbook element from the namespace
// http://schemas.openxmlformats.org/spreadsheetml/2006/main
type xlsxWorkbook struct {
	Sheets xlsxSheets `xml:"sheets"`
}

// xlsxSheets directly maps the sheets element from the namespace
// http://schemas.openxmlformats.org/spreadsheetml/2006/main
type xlsxSheets struct {
	Sheet []xlsxSheet `xml:"sheet"`
}

// xlsxSheet directly maps the sheet element from the namespace
// http://schemas.openxmlformats.org/spreadsheetml/2006/main
type xlsxSheet struct {
	Name    string `xml:"name,attr,omitempty"`
	SheetID string `xml:"sheetId,attr,omitempty"`
}
<commit_msg>简化readWorkbook<commit_after>package internal

import (
	"encoding/xml"
	"io"
)

// Ref: tealeg/xlsx/xmlWorkbook.go
// I just keep the part I want.

func readWorkbookXML(rd io.Reader) (*xlsxWorkbook, error) {
	var err error
	workbook := new(xlsxWorkbook)
	decoder := xml.NewDecoder(rd)
	err = decoder.Decode(workbook)
	if err != nil {
		return nil, err
	}
	return workbook, nil
}

// xlsxWorkbook directly maps the workbook element from the namespace
// http://schemas.openxmlformats.org/spreadsheetml/2006/main
type xlsxWorkbook struct {
	Sheets xlsxSheets `xml:"sheets"`
}

// xlsxSheets directly maps the sheets element from the namespace
// http://schemas.openxmlformats.org/spreadsheetml/2006/main
type xlsxSheets struct {
	Sheet []xlsxSheet `xml:"sheet"`
}

// xlsxSheet directly maps the sheet element from the namespace
// http://schemas.openxmlformats.org/spreadsheetml/2006/main
type xlsxSheet struct {
	Name    string `xml:"name,attr,omitempty"`
	SheetID string `xml:"sheetId,attr,omitempty"`
}
<|endoftext|>


<commit_before>//
// Copyright (c) 2016 Intel Corporation
//
// 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 hyperstart

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"math"
	"net"
	"sync"
	"time"

	hyper "github.com/hyperhq/runv/hyperstart/api/json"
)

// Control command IDs
// Need to be in sync with hyperstart/src/api.h
const (
	Version         = "version"
	StartPod        = "startpod"
	DestroyPod      = "destroypod"
	ExecCmd         = "execcmd"
	Ready           = "ready"
	Ack             = "ack"
	Error           = "error"
	WinSize         = "winsize"
	Ping            = "ping"
	FinishPod       = "finishpod"
	Next            = "next"
	WriteFile       = "writefile"
	ReadFile        = "readfile"
	NewContainer    = "newcontainer"
	KillContainer   = "killcontainer"
	RemoveContainer = "removecontainer"
	OnlineCPUMem    = "onlinecpumem"
	SetupInterface  = "setupinterface"
	SetupRoute      = "setuproute"
)

var codeList = map[string]uint32{
	Version:         hyper.INIT_VERSION,
	StartPod:        hyper.INIT_STARTPOD,
	DestroyPod:      hyper.INIT_DESTROYPOD,
	ExecCmd:         hyper.INIT_EXECCMD,
	Ready:           hyper.INIT_READY,
	Ack:             hyper.INIT_ACK,
	Error:           hyper.INIT_ERROR,
	WinSize:         hyper.INIT_WINSIZE,
	Ping:            hyper.INIT_PING,
	Next:            hyper.INIT_NEXT,
	WriteFile:       hyper.INIT_WRITEFILE,
	ReadFile:        hyper.INIT_READFILE,
	NewContainer:    hyper.INIT_NEWCONTAINER,
	KillContainer:   hyper.INIT_KILLCONTAINER,
	RemoveContainer: hyper.INIT_REMOVECONTAINER,
	OnlineCPUMem:    hyper.INIT_ONLINECPUMEM,
	SetupInterface:  hyper.INIT_SETUPINTERFACE,
	SetupRoute:      hyper.INIT_SETUPROUTE,
}

// Values related to the communication on control channel.
const (
	ctlHdrSize      = 8
	ctlHdrLenOffset = 4
)

// Values related to the communication on tty channel.
const (
	ttyHdrSize      = 12
	ttyHdrLenOffset = 8
)

type connState struct {
	sync.Mutex
	opened bool
}

func (c *connState) close() {
	c.Lock()
	defer c.Unlock()

	c.opened = false
}

func (c *connState) open() {
	c.Lock()
	defer c.Unlock()

	c.opened = true
}

func (c *connState) closed() bool {
	c.Lock()
	defer c.Unlock()

	return !c.opened
}

// Hyperstart is the base structure for hyperstart.
type Hyperstart struct {
	ctlSerial, ioSerial string
	sockType            string
	ctl, io             net.Conn
	ctlState, ioState   connState

	// ctl access is arbitrated by ctlMutex. We can only allow a single
	// "transaction" (write command + read answer) at a time
	ctlMutex sync.Mutex
}

// NewHyperstart returns a new hyperstart structure.
func NewHyperstart(ctlSerial, ioSerial, sockType string) *Hyperstart {
	return &Hyperstart{
		ctlSerial: ctlSerial,
		ioSerial:  ioSerial,
		sockType:  sockType,
	}
}

// OpenSockets opens both CTL and IO sockets.
func (h *Hyperstart) OpenSockets() error {
	var err error

	h.ctl, err = net.Dial(h.sockType, h.ctlSerial)
	if err != nil {
		return err
	}
	h.ctlState.open()

	h.io, err = net.Dial(h.sockType, h.ioSerial)
	if err != nil {
		h.ctl.Close()
		return err
	}
	h.ioState.open()

	return nil
}

// CloseSockets closes both CTL and IO sockets.
func (h *Hyperstart) CloseSockets() error {
	if !h.ctlState.closed() {
		err := h.ctl.Close()
		if err != nil {
			return err
		}

		h.ctlState.close()
	}

	if !h.ioState.closed() {
		err := h.io.Close()
		if err != nil {
			return err
		}

		h.ioState.close()
	}

	return nil
}

// SetDeadline sets a timeout for CTL connection.
func (h *Hyperstart) SetDeadline(t time.Time) error {
	err := h.ctl.SetDeadline(t)
	if err != nil {
		return err
	}

	return nil
}

// IsStarted returns about connection status.
func (h *Hyperstart) IsStarted() bool {
	ret := false
	timeoutDuration := 1 * time.Second

	if h.ctlState.closed() {
		return ret
	}

	h.SetDeadline(time.Now().Add(timeoutDuration))

	_, err := h.SendCtlMessage(Ping, nil)
	if err == nil {
		ret = true
	}

	h.SetDeadline(time.Time{})

	if ret == false {
		h.CloseSockets()
	}

	return ret
}

// FormatMessage formats hyperstart messages.
func FormatMessage(payload interface{}) ([]byte, error) {
	var payloadSlice []byte
	var err error

	if payload != nil {
		switch p := payload.(type) {
		case string:
			payloadSlice = []byte(p)
		default:
			payloadSlice, err = json.Marshal(p)
			if err != nil {
				return nil, err
			}
		}
	}

	return payloadSlice, nil
}

func (h *Hyperstart) readCtlMessage() (*hyper.DecodedMessage, error) {
	needRead := ctlHdrSize
	length := 0
	read := 0
	buf := make([]byte, 512)
	res := []byte{}
	for read < needRead {
		want := needRead - read
		if want > 512 {
			want = 512
		}
		nr, err := h.ctl.Read(buf[:want])
		if err != nil {
			return nil, err
		}

		res = append(res, buf[:nr]...)
		read = read + nr

		if length == 0 && read >= ctlHdrSize {
			length = int(binary.BigEndian.Uint32(res[ctlHdrLenOffset:ctlHdrSize]))
			if length > ctlHdrSize {
				needRead = length
			}
		}
	}

	return &hyper.DecodedMessage{
		Code:    binary.BigEndian.Uint32(res[:ctlHdrLenOffset]),
		Message: res[ctlHdrSize:],
	}, nil
}

func (h *Hyperstart) writeCtlMessage(m *hyper.DecodedMessage) error {
	length := len(m.Message) + ctlHdrSize
	// XXX: Support sending messages by chunks to support messages over
	// 10240 bytes. That limit is from hyperstart src/init.c,
	// hyper_channel_ops, rbuf_size.
	if length > 10240 {
		return fmt.Errorf("message too long %d", length)
	}
	msg := make([]byte, length)
	binary.BigEndian.PutUint32(msg[:], uint32(m.Code))
	binary.BigEndian.PutUint32(msg[ctlHdrLenOffset:], uint32(length))
	copy(msg[ctlHdrSize:], m.Message)

	_, err := h.ctl.Write(msg)
	if err != nil {
		return err
	}

	return nil
}

// ReadIoMessageWithConn returns data coming from the specified IO channel.
func ReadIoMessageWithConn(conn net.Conn) (*hyper.TtyMessage, error) {
	needRead := ttyHdrSize
	length := 0
	read := 0
	buf := make([]byte, 512)
	res := []byte{}
	for read < needRead {
		want := needRead - read
		if want > 512 {
			want = 512
		}
		nr, err := conn.Read(buf[:want])
		if err != nil {
			return nil, err
		}

		res = append(res, buf[:nr]...)
		read = read + nr

		if length == 0 && read >= ttyHdrSize {
			length = int(binary.BigEndian.Uint32(res[ttyHdrLenOffset:ttyHdrSize]))
			if length > ttyHdrSize {
				needRead = length
			}
		}
	}

	return &hyper.TtyMessage{
		Session: binary.BigEndian.Uint64(res[:ttyHdrLenOffset]),
		Message: res[ttyHdrSize:],
	}, nil
}

// ReadIoMessage returns data coming from the IO channel.
func (h *Hyperstart) ReadIoMessage() (*hyper.TtyMessage, error) {
	return ReadIoMessageWithConn(h.io)
}

// SendIoMessageWithConn sends data to the specified IO channel.
func SendIoMessageWithConn(conn net.Conn, ttyMsg *hyper.TtyMessage) error {
	length := len(ttyMsg.Message) + ttyHdrSize
	// XXX: Support sending messages by chunks to support messages over
	// 10240 bytes. That limit is from hyperstart src/init.c,
	// hyper_channel_ops, rbuf_size.
	if length > 10240 {
		return fmt.Errorf("message too long %d", length)
	}
	msg := make([]byte, length)
	binary.BigEndian.PutUint64(msg[:], ttyMsg.Session)
	binary.BigEndian.PutUint32(msg[ttyHdrLenOffset:], uint32(length))
	copy(msg[ttyHdrSize:], ttyMsg.Message)

	n, err := conn.Write(msg)
	if err != nil {
		return err
	}

	if n != length {
		return fmt.Errorf("%d bytes written out of %d expected", n, length)
	}

	return nil
}

// SendIoMessage sends data to the IO channel.
func (h *Hyperstart) SendIoMessage(ttyMsg *hyper.TtyMessage) error {
	return SendIoMessageWithConn(h.io, ttyMsg)
}

func codeFromCmd(cmd string) (uint32, error) {
	_, ok := codeList[cmd]
	if ok == false {
		return math.MaxUint32, fmt.Errorf("unknown command '%s'", cmd)
	}

	return codeList[cmd], nil
}

func (h *Hyperstart) expectReadingCmd(code uint32) (*hyper.DecodedMessage, error) {
	var msg *hyper.DecodedMessage
	var err error

	for {
		msg, err = h.readCtlMessage()
		if err != nil {
			return nil, nil
		}

		if msg.Code == code {
			break
		}

		if msg.Code == hyper.INIT_NEXT || msg.Code == hyper.INIT_READY {
			continue
		}

		if msg.Code != code {
			if msg.Code == hyper.INIT_ERROR {
				return nil, fmt.Errorf("ERROR received from Hyperstart")
			}

			return nil, fmt.Errorf("CMD ID received %d not matching expected %d", msg.Code, code)
		}
	}

	return msg, nil
}

// WaitForReady waits for a READY message on CTL channel.
func (h *Hyperstart) WaitForReady() error {
	h.ctlMutex.Lock()
	defer h.ctlMutex.Unlock()

	_, err := h.expectReadingCmd(hyper.INIT_READY)
	if err != nil {
		return err
	}

	return nil
}

// SendCtlMessage sends a message to the CTL channel.
func (h *Hyperstart) SendCtlMessage(cmd string, data []byte) (*hyper.DecodedMessage, error) {
	h.ctlMutex.Lock()
	defer h.ctlMutex.Unlock()

	code, err := codeFromCmd(cmd)
	if err != nil {
		return nil, err
	}

	// Write message
	msg := &hyper.DecodedMessage{
		Code:    code,
		Message: data,
	}
	err = h.writeCtlMessage(msg)
	if err != nil {
		return nil, err
	}

	// Wait for answer
	resp, err := h.expectReadingCmd(hyper.INIT_ACK)
	if err != nil {
		return nil, err
	}

	return resp, nil
}
<commit_msg>hyperstart: Add some precisions about SendCtlMessage<commit_after>//
// Copyright (c) 2016 Intel Corporation
//
// 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 hyperstart

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"math"
	"net"
	"sync"
	"time"

	hyper "github.com/hyperhq/runv/hyperstart/api/json"
)

// Control command IDs
// Need to be in sync with hyperstart/src/api.h
const (
	Version         = "version"
	StartPod        = "startpod"
	DestroyPod      = "destroypod"
	ExecCmd         = "execcmd"
	Ready           = "ready"
	Ack             = "ack"
	Error           = "error"
	WinSize         = "winsize"
	Ping            = "ping"
	FinishPod       = "finishpod"
	Next            = "next"
	WriteFile       = "writefile"
	ReadFile        = "readfile"
	NewContainer    = "newcontainer"
	KillContainer   = "killcontainer"
	RemoveContainer = "removecontainer"
	OnlineCPUMem    = "onlinecpumem"
	SetupInterface  = "setupinterface"
	SetupRoute      = "setuproute"
)

var codeList = map[string]uint32{
	Version:         hyper.INIT_VERSION,
	StartPod:        hyper.INIT_STARTPOD,
	DestroyPod:      hyper.INIT_DESTROYPOD,
	ExecCmd:         hyper.INIT_EXECCMD,
	Ready:           hyper.INIT_READY,
	Ack:             hyper.INIT_ACK,
	Error:           hyper.INIT_ERROR,
	WinSize:         hyper.INIT_WINSIZE,
	Ping:            hyper.INIT_PING,
	Next:            hyper.INIT_NEXT,
	WriteFile:       hyper.INIT_WRITEFILE,
	ReadFile:        hyper.INIT_READFILE,
	NewContainer:    hyper.INIT_NEWCONTAINER,
	KillContainer:   hyper.INIT_KILLCONTAINER,
	RemoveContainer: hyper.INIT_REMOVECONTAINER,
	OnlineCPUMem:    hyper.INIT_ONLINECPUMEM,
	SetupInterface:  hyper.INIT_SETUPINTERFACE,
	SetupRoute:      hyper.INIT_SETUPROUTE,
}

// Values related to the communication on control channel.
const (
	ctlHdrSize      = 8
	ctlHdrLenOffset = 4
)

// Values related to the communication on tty channel.
const (
	ttyHdrSize      = 12
	ttyHdrLenOffset = 8
)

type connState struct {
	sync.Mutex
	opened bool
}

func (c *connState) close() {
	c.Lock()
	defer c.Unlock()

	c.opened = false
}

func (c *connState) open() {
	c.Lock()
	defer c.Unlock()

	c.opened = true
}

func (c *connState) closed() bool {
	c.Lock()
	defer c.Unlock()

	return !c.opened
}

// Hyperstart is the base structure for hyperstart.
type Hyperstart struct {
	ctlSerial, ioSerial string
	sockType            string
	ctl, io             net.Conn
	ctlState, ioState   connState

	// ctl access is arbitrated by ctlMutex. We can only allow a single
	// "transaction" (write command + read answer) at a time
	ctlMutex sync.Mutex
}

// NewHyperstart returns a new hyperstart structure.
func NewHyperstart(ctlSerial, ioSerial, sockType string) *Hyperstart {
	return &Hyperstart{
		ctlSerial: ctlSerial,
		ioSerial:  ioSerial,
		sockType:  sockType,
	}
}

// OpenSockets opens both CTL and IO sockets.
func (h *Hyperstart) OpenSockets() error {
	var err error

	h.ctl, err = net.Dial(h.sockType, h.ctlSerial)
	if err != nil {
		return err
	}
	h.ctlState.open()

	h.io, err = net.Dial(h.sockType, h.ioSerial)
	if err != nil {
		h.ctl.Close()
		return err
	}
	h.ioState.open()

	return nil
}

// CloseSockets closes both CTL and IO sockets.
func (h *Hyperstart) CloseSockets() error {
	if !h.ctlState.closed() {
		err := h.ctl.Close()
		if err != nil {
			return err
		}

		h.ctlState.close()
	}

	if !h.ioState.closed() {
		err := h.io.Close()
		if err != nil {
			return err
		}

		h.ioState.close()
	}

	return nil
}

// SetDeadline sets a timeout for CTL connection.
func (h *Hyperstart) SetDeadline(t time.Time) error {
	err := h.ctl.SetDeadline(t)
	if err != nil {
		return err
	}

	return nil
}

// IsStarted returns about connection status.
func (h *Hyperstart) IsStarted() bool {
	ret := false
	timeoutDuration := 1 * time.Second

	if h.ctlState.closed() {
		return ret
	}

	h.SetDeadline(time.Now().Add(timeoutDuration))

	_, err := h.SendCtlMessage(Ping, nil)
	if err == nil {
		ret = true
	}

	h.SetDeadline(time.Time{})

	if ret == false {
		h.CloseSockets()
	}

	return ret
}

// FormatMessage formats hyperstart messages.
func FormatMessage(payload interface{}) ([]byte, error) {
	var payloadSlice []byte
	var err error

	if payload != nil {
		switch p := payload.(type) {
		case string:
			payloadSlice = []byte(p)
		default:
			payloadSlice, err = json.Marshal(p)
			if err != nil {
				return nil, err
			}
		}
	}

	return payloadSlice, nil
}

func (h *Hyperstart) readCtlMessage() (*hyper.DecodedMessage, error) {
	needRead := ctlHdrSize
	length := 0
	read := 0
	buf := make([]byte, 512)
	res := []byte{}
	for read < needRead {
		want := needRead - read
		if want > 512 {
			want = 512
		}
		nr, err := h.ctl.Read(buf[:want])
		if err != nil {
			return nil, err
		}

		res = append(res, buf[:nr]...)
		read = read + nr

		if length == 0 && read >= ctlHdrSize {
			length = int(binary.BigEndian.Uint32(res[ctlHdrLenOffset:ctlHdrSize]))
			if length > ctlHdrSize {
				needRead = length
			}
		}
	}

	return &hyper.DecodedMessage{
		Code:    binary.BigEndian.Uint32(res[:ctlHdrLenOffset]),
		Message: res[ctlHdrSize:],
	}, nil
}

func (h *Hyperstart) writeCtlMessage(m *hyper.DecodedMessage) error {
	length := len(m.Message) + ctlHdrSize
	// XXX: Support sending messages by chunks to support messages over
	// 10240 bytes. That limit is from hyperstart src/init.c,
	// hyper_channel_ops, rbuf_size.
	if length > 10240 {
		return fmt.Errorf("message too long %d", length)
	}
	msg := make([]byte, length)
	binary.BigEndian.PutUint32(msg[:], uint32(m.Code))
	binary.BigEndian.PutUint32(msg[ctlHdrLenOffset:], uint32(length))
	copy(msg[ctlHdrSize:], m.Message)

	_, err := h.ctl.Write(msg)
	if err != nil {
		return err
	}

	return nil
}

// ReadIoMessageWithConn returns data coming from the specified IO channel.
func ReadIoMessageWithConn(conn net.Conn) (*hyper.TtyMessage, error) {
	needRead := ttyHdrSize
	length := 0
	read := 0
	buf := make([]byte, 512)
	res := []byte{}
	for read < needRead {
		want := needRead - read
		if want > 512 {
			want = 512
		}
		nr, err := conn.Read(buf[:want])
		if err != nil {
			return nil, err
		}

		res = append(res, buf[:nr]...)
		read = read + nr

		if length == 0 && read >= ttyHdrSize {
			length = int(binary.BigEndian.Uint32(res[ttyHdrLenOffset:ttyHdrSize]))
			if length > ttyHdrSize {
				needRead = length
			}
		}
	}

	return &hyper.TtyMessage{
		Session: binary.BigEndian.Uint64(res[:ttyHdrLenOffset]),
		Message: res[ttyHdrSize:],
	}, nil
}

// ReadIoMessage returns data coming from the IO channel.
func (h *Hyperstart) ReadIoMessage() (*hyper.TtyMessage, error) {
	return ReadIoMessageWithConn(h.io)
}

// SendIoMessageWithConn sends data to the specified IO channel.
func SendIoMessageWithConn(conn net.Conn, ttyMsg *hyper.TtyMessage) error {
	length := len(ttyMsg.Message) + ttyHdrSize
	// XXX: Support sending messages by chunks to support messages over
	// 10240 bytes. That limit is from hyperstart src/init.c,
	// hyper_channel_ops, rbuf_size.
	if length > 10240 {
		return fmt.Errorf("message too long %d", length)
	}
	msg := make([]byte, length)
	binary.BigEndian.PutUint64(msg[:], ttyMsg.Session)
	binary.BigEndian.PutUint32(msg[ttyHdrLenOffset:], uint32(length))
	copy(msg[ttyHdrSize:], ttyMsg.Message)

	n, err := conn.Write(msg)
	if err != nil {
		return err
	}

	if n != length {
		return fmt.Errorf("%d bytes written out of %d expected", n, length)
	}

	return nil
}

// SendIoMessage sends data to the IO channel.
func (h *Hyperstart) SendIoMessage(ttyMsg *hyper.TtyMessage) error {
	return SendIoMessageWithConn(h.io, ttyMsg)
}

func codeFromCmd(cmd string) (uint32, error) {
	_, ok := codeList[cmd]
	if ok == false {
		return math.MaxUint32, fmt.Errorf("unknown command '%s'", cmd)
	}

	return codeList[cmd], nil
}

func (h *Hyperstart) expectReadingCmd(code uint32) (*hyper.DecodedMessage, error) {
	var msg *hyper.DecodedMessage
	var err error

	for {
		msg, err = h.readCtlMessage()
		if err != nil {
			return nil, nil
		}

		if msg.Code == code {
			break
		}

		if msg.Code == hyper.INIT_NEXT || msg.Code == hyper.INIT_READY {
			continue
		}

		if msg.Code != code {
			if msg.Code == hyper.INIT_ERROR {
				return nil, fmt.Errorf("ERROR received from Hyperstart")
			}

			return nil, fmt.Errorf("CMD ID received %d not matching expected %d", msg.Code, code)
		}
	}

	return msg, nil
}

// WaitForReady waits for a READY message on CTL channel.
func (h *Hyperstart) WaitForReady() error {
	h.ctlMutex.Lock()
	defer h.ctlMutex.Unlock()

	_, err := h.expectReadingCmd(hyper.INIT_READY)
	if err != nil {
		return err
	}

	return nil
}

// SendCtlMessage sends a message to the CTL channel.
//
// This function does a full transaction over the CTL channel: it will write a
// command and wait for hyperstart's answer before returning. Several
// concurrent calls to SendCtlMessage is allowed, the function ensuring proper
// serialization of the communication over the underlying hyperstart serial
// link.
func (h *Hyperstart) SendCtlMessage(cmd string, data []byte) (*hyper.DecodedMessage, error) {
	h.ctlMutex.Lock()
	defer h.ctlMutex.Unlock()

	code, err := codeFromCmd(cmd)
	if err != nil {
		return nil, err
	}

	// Write message
	msg := &hyper.DecodedMessage{
		Code:    code,
		Message: data,
	}
	err = h.writeCtlMessage(msg)
	if err != nil {
		return nil, err
	}

	// Wait for answer
	resp, err := h.expectReadingCmd(hyper.INIT_ACK)
	if err != nil {
		return nil, err
	}

	return resp, nil
}
<|endoftext|>


<commit_before>package messages

import (
	"fmt"
	"gnd.la/log"
	"gnd.la/util/astutil"
	"gnd.la/util/pkgutil"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"strings"
)

func DefaultFunctions() []*Function {
	return []*Function{
		// Singular functions without context
		{Name: "gnd.la/i18n.T"},
		{Name: "gnd.la/i18n.Errorf"},
		{Name: "gnd.la/i18n.Sprintf"},
		{Name: "gnd.la/i18n.NewError"},
		{Name: "gnd.la/mux.Context.T"},
		{Name: "T", Template: true},
		// Singular functions with context
		{Name: "gnd.la/i18n.Tc", Context: true},
		{Name: "gnd.la/i18n.Sprintfc", Context: true},
		{Name: "gnd.la/i18n.Errorfc", Context: true},
		{Name: "gnd.la/i18n.NewErrorc", Context: true},
		{Name: "gnd.la/mux.Context.Tc", Context: true},
		{Name: "Tc", Template: true, Context: true},
		// Plural functions without context
		{Name: "gnd.la/i18n.Tn", Plural: true},
		{Name: "gnd.la/i18n.Sprintfn", Plural: true},
		{Name: "gnd.la/i18n.Errorfn", Plural: true},
		{Name: "gnd.la/i18n.NewErrorn", Plural: true},
		{Name: "gnd.la/mux.Context.Tn", Plural: true},
		{Name: "Tn", Template: true, Plural: true},
		// Plural functions with context
		{Name: "gnd.la/i18n.Tnc", Context: true, Plural: true},
		{Name: "gnd.la/i18n.Errorfnc", Context: true, Plural: true},
		{Name: "gnd.la/i18n.Sprintfnc", Context: true, Plural: true},
		{Name: "gnd.la/i18n.NewErrornc", Context: true, Plural: true},
		{Name: "gnd.la/mux.Context.Tnc", Context: true, Plural: true},
		{Name: "Tnc", Template: true, Context: true, Plural: true},
	}
}

func DefaultTypes() []string {
	return []string{
		"gnd.la/i18n.String",
	}
}

func DefaultTagFields() []string {
	return []string{
		"help",
		"label",
		"placeholder",
	}
}

func Extract(dir string, functions []*Function, types []string, tagFields []string) ([]*Message, error) {
	messages := make(messageMap)
	err := extract(messages, dir, functions, types, tagFields)
	if err != nil {
		return nil, err
	}
	return messages.Messages(), nil
}

func extract(messages messageMap, dir string, functions []*Function, types []string, tagFields []string) error {
	f, err := os.Open(dir)
	if err != nil {
		return err
	}
	defer f.Close()
	infos, err := f.Readdir(-1)
	if err != nil {
		return err
	}
	for _, v := range infos {
		name := v.Name()
		p := filepath.Join(dir, name)
		if v.IsDir() {
			if !pkgutil.IsPackage(p) {
				if err := extract(messages, p, functions, types, tagFields); err != nil {
					return err
				}
			}
			continue
		}
		switch strings.ToLower(filepath.Ext(name)) {
		// TODO: templates, strings files
		case ".go":
			if err := extractGoMessages(messages, p, functions, types, tagFields); err != nil {
				return err
			}
		case ".po", ".pot":
			// Do nothing
		}
	}
	return nil
}

func extractGoMessages(messages messageMap, path string, functions []*Function, types []string, tagFields []string) error {
	log.Debugf("Extracting messages from Go file %s", path)
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
	if err != nil {
		return fmt.Errorf("error parsing go file %s: %s", path, err)
	}
	for _, v := range functions {
		if v.Template {
			continue
		}
		if err := extractGoFunc(messages, fset, f, v); err != nil {
			return err
		}
	}
	for _, v := range types {
		if err := extractGoType(messages, fset, f, v); err != nil {
			return err
		}
	}
	for _, v := range tagFields {
		if err := extractGoTagField(messages, fset, f, v); err != nil {
			return err
		}
	}
	return nil
}

func extractGoFunc(messages messageMap, fset *token.FileSet, f *ast.File, fn *Function) error {
	calls, err := astutil.Calls(fset, f, fn.Name)
	if err != nil {
		return err
	}
	n := 0
	if fn.Context {
		n++
	}
	var message *Message
	var position *token.Position
	for _, c := range calls {
		if fn.Plural {
			if len(c.Args) < n+3 {
				log.Debugf("Skipping plural function %s (%v) - not enough arguments", astutil.Ident(c.Fun), fset.Position(c.Pos()))
				continue
			}
			slit, spos := astutil.StringLiteral(fset, c.Args[n])
			if slit == "" || spos == nil {
				log.Debugf("Skipping first argument to plural function %s (%v) - not a literal", astutil.Ident(c.Fun), fset.Position(c.Pos()))
				continue
			}
			plit, ppos := astutil.StringLiteral(fset, c.Args[n+1])
			if plit == "" || ppos == nil {
				log.Debugf("Skipping second argument to plural function %s (%v) - not a literal", astutil.Ident(c.Fun), fset.Position(c.Pos()))
				continue
			}
			message = &Message{
				Singular: slit,
				Plural:   plit,
			}
			position = spos
		} else {
			if len(c.Args) < n+1 {
				log.Debugf("Skipping singular function %s (%v) - not enough arguments", astutil.Ident(c.Fun), fset.Position(c.Pos()))
				continue
			}
			lit, pos := astutil.StringLiteral(fset, c.Args[n])
			if lit == "" || pos == nil {
				log.Debugf("Skipping argument to singular function %s (%v) - not a literal", astutil.Ident(c.Fun), fset.Position(c.Pos()))
				continue
			}
			message = &Message{
				Singular: lit,
			}
			position = pos
		}
		if message != nil && position != nil {
			if fn.Context {
				ctx, cpos := astutil.StringLiteral(fset, c.Args[0])
				if ctx == "" || cpos == nil {
					log.Debugf("Skipping argument to context function %s (%v) - empty context", astutil.Ident(c.Fun), fset.Position(c.Pos()))
					continue
				}
				message.Context = ctx
			}
			if err := messages.Add(message, position, comments(fset, f, position)); err != nil {
				return err
			}
		}
	}
	return nil
}

func extractGoType(messages messageMap, fset *token.FileSet, f *ast.File, typ string) error {
	// for castings
	tf := &Function{Name: typ}
	if err := extractGoFunc(messages, fset, f, tf); err != nil {
		return err
	}
	strings, err := astutil.Strings(fset, f, typ)
	if err != nil {
		return err
	}
	for _, s := range strings {
		comment := comments(fset, f, s.Position)
		if err := messages.AddString(s, comment); err != nil {
			return err
		}
	}
	return nil
}

func extractGoTagField(messages messageMap, fset *token.FileSet, f *ast.File, tagField string) error {
	strings, err := astutil.TagFields(fset, f, tagField)
	if err != nil {
		return err
	}
	for _, s := range strings {
		comment := comments(fset, f, s.Position)
		if err := messages.AddString(s, comment); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>Add support for extracting strings from templates<commit_after>package messages

import (
	"fmt"
	"gnd.la/log"
	"gnd.la/util/astutil"
	"gnd.la/util/pkgutil"
	"gnd.la/util/templateutil"
	"go/ast"
	"go/parser"
	"go/token"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"text/template/parse"
)

func DefaultFunctions() []*Function {
	return []*Function{
		// Singular functions without context
		{Name: "gnd.la/i18n.T"},
		{Name: "gnd.la/i18n.Errorf"},
		{Name: "gnd.la/i18n.Sprintf"},
		{Name: "gnd.la/i18n.NewError"},
		{Name: "gnd.la/mux.Context.T"},
		{Name: "T", Template: true},
		// Singular functions with context
		{Name: "gnd.la/i18n.Tc", Context: true},
		{Name: "gnd.la/i18n.Sprintfc", Context: true},
		{Name: "gnd.la/i18n.Errorfc", Context: true},
		{Name: "gnd.la/i18n.NewErrorc", Context: true},
		{Name: "gnd.la/mux.Context.Tc", Context: true},
		{Name: "Tc", Template: true, Context: true},
		// Plural functions without context
		{Name: "gnd.la/i18n.Tn", Plural: true},
		{Name: "gnd.la/i18n.Sprintfn", Plural: true},
		{Name: "gnd.la/i18n.Errorfn", Plural: true},
		{Name: "gnd.la/i18n.NewErrorn", Plural: true},
		{Name: "gnd.la/mux.Context.Tn", Plural: true},
		{Name: "Tn", Template: true, Plural: true},
		// Plural functions with context
		{Name: "gnd.la/i18n.Tnc", Context: true, Plural: true},
		{Name: "gnd.la/i18n.Errorfnc", Context: true, Plural: true},
		{Name: "gnd.la/i18n.Sprintfnc", Context: true, Plural: true},
		{Name: "gnd.la/i18n.NewErrornc", Context: true, Plural: true},
		{Name: "gnd.la/mux.Context.Tnc", Context: true, Plural: true},
		{Name: "Tnc", Template: true, Context: true, Plural: true},
	}
}

func DefaultTypes() []string {
	return []string{
		"gnd.la/i18n.String",
	}
}

func DefaultTagFields() []string {
	return []string{
		"help",
		"label",
		"placeholder",
	}
}

func Extract(dir string, functions []*Function, types []string, tagFields []string) ([]*Message, error) {
	messages := make(messageMap)
	err := extract(messages, dir, functions, types, tagFields)
	if err != nil {
		return nil, err
	}
	return messages.Messages(), nil
}

func extract(messages messageMap, dir string, functions []*Function, types []string, tagFields []string) error {
	f, err := os.Open(dir)
	if err != nil {
		return err
	}
	defer f.Close()
	infos, err := f.Readdir(-1)
	if err != nil {
		return err
	}
	for _, v := range infos {
		name := v.Name()
		p := filepath.Join(dir, name)
		if v.IsDir() {
			if !pkgutil.IsPackage(p) {
				if err := extract(messages, p, functions, types, tagFields); err != nil {
					return err
				}
			}
			continue
		}
		switch strings.ToLower(filepath.Ext(name)) {
		// TODO: text and strings files
		case ".html":
			if err := extractTemplateMessages(messages, p, functions); err != nil {
				return err
			}
		case ".go":
			if err := extractGoMessages(messages, p, functions, types, tagFields); err != nil {
				return err
			}
		case ".po", ".pot":
			// Do nothing
		}
	}
	return nil
}

func extractGoMessages(messages messageMap, path string, functions []*Function, types []string, tagFields []string) error {
	log.Debugf("Extracting messages from Go file %s", path)
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
	if err != nil {
		return fmt.Errorf("error parsing go file %s: %s", path, err)
	}
	for _, v := range functions {
		if v.Template {
			continue
		}
		if err := extractGoFunc(messages, fset, f, v); err != nil {
			return err
		}
	}
	for _, v := range types {
		if err := extractGoType(messages, fset, f, v); err != nil {
			return err
		}
	}
	for _, v := range tagFields {
		if err := extractGoTagField(messages, fset, f, v); err != nil {
			return err
		}
	}
	return nil
}

func extractGoFunc(messages messageMap, fset *token.FileSet, f *ast.File, fn *Function) error {
	calls, err := astutil.Calls(fset, f, fn.Name)
	if err != nil {
		return err
	}
	n := 0
	if fn.Context {
		n++
	}
	var message *Message
	var position *token.Position
	for _, c := range calls {
		if fn.Plural {
			if len(c.Args) < n+3 {
				log.Debugf("Skipping plural function %s (%v) - not enough arguments", astutil.Ident(c.Fun), fset.Position(c.Pos()))
				continue
			}
			slit, spos := astutil.StringLiteral(fset, c.Args[n])
			if slit == "" || spos == nil {
				log.Debugf("Skipping first argument to plural function %s (%v) - not a literal", astutil.Ident(c.Fun), fset.Position(c.Pos()))
				continue
			}
			plit, ppos := astutil.StringLiteral(fset, c.Args[n+1])
			if plit == "" || ppos == nil {
				log.Debugf("Skipping second argument to plural function %s (%v) - not a literal", astutil.Ident(c.Fun), fset.Position(c.Pos()))
				continue
			}
			message = &Message{
				Singular: slit,
				Plural:   plit,
			}
			position = spos
		} else {
			if len(c.Args) < n+1 {
				log.Debugf("Skipping singular function %s (%v) - not enough arguments", astutil.Ident(c.Fun), fset.Position(c.Pos()))
				continue
			}
			lit, pos := astutil.StringLiteral(fset, c.Args[n])
			if lit == "" || pos == nil {
				log.Debugf("Skipping argument to singular function %s (%v) - not a literal", astutil.Ident(c.Fun), fset.Position(c.Pos()))
				continue
			}
			message = &Message{
				Singular: lit,
			}
			position = pos
		}
		if message != nil && position != nil {
			if fn.Context {
				ctx, cpos := astutil.StringLiteral(fset, c.Args[0])
				if ctx == "" || cpos == nil {
					log.Debugf("Skipping argument to context function %s (%v) - empty context", astutil.Ident(c.Fun), fset.Position(c.Pos()))
					continue
				}
				message.Context = ctx
			}
			if err := messages.Add(message, position, comments(fset, f, position)); err != nil {
				return err
			}
		}
	}
	return nil
}

func extractGoType(messages messageMap, fset *token.FileSet, f *ast.File, typ string) error {
	// for castings
	tf := &Function{Name: typ}
	if err := extractGoFunc(messages, fset, f, tf); err != nil {
		return err
	}
	strings, err := astutil.Strings(fset, f, typ)
	if err != nil {
		return err
	}
	for _, s := range strings {
		comment := comments(fset, f, s.Position)
		if err := messages.AddString(s, comment); err != nil {
			return err
		}
	}
	return nil
}

func extractGoTagField(messages messageMap, fset *token.FileSet, f *ast.File, tagField string) error {
	strings, err := astutil.TagFields(fset, f, tagField)
	if err != nil {
		return err
	}
	for _, s := range strings {
		comment := comments(fset, f, s.Position)
		if err := messages.AddString(s, comment); err != nil {
			return err
		}
	}
	return nil
}

func extractTemplateMessages(messages messageMap, path string, functions []*Function) error {
	log.Debugf("Extracting messages from template file %s", path)
	b, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}
	funcs := make(map[string]*Function)
	for _, v := range functions {
		if v.Template {
			funcs[v.Name] = v
		}
	}
	text := string(b)
	treeSet, err := templateutil.Parse(path, text)
	if err != nil {
		return err
	}
	for _, v := range treeSet {
		templateutil.WalkTree(v, func(n, p parse.Node) {
			var fname string
			switch n.Type() {
			case parse.NodeIdentifier:
				fname = n.(*parse.IdentifierNode).Ident
			case parse.NodeField:
				ident := n.(*parse.FieldNode).Ident
				if len(ident) > 1 {
					fname = ident[len(ident)-1]
				}
			case parse.NodeVariable:
				ident := n.(*parse.VariableNode).Ident
				if len(ident) > 1 {
					fname = ident[len(ident)-1]
				}
			}
			if fname != "" {
				f := funcs[fname]
				if f != nil {
					count := 1
					if f.Context {
						count++
					}
					if f.Plural {
						count++
					}
					cmd := p.(*parse.CommandNode)
					// First argument is the function name
					if c := len(cmd.Args) - 1; c != count {
						log.Debugf("Skipping function %s (%v) - want %d arguments, got %d", f.Name, n.Position(), count, c)
						return
					}
					var s []string
					for ii := 1; ii < len(cmd.Args); ii++ {
						if sn, ok := cmd.Args[ii].(*parse.StringNode); ok {
							s = append(s, sn.Text)
						} else {
							log.Debugf("Skipping function %s (%v) - non-string argument at position %d", f.Name, n.Position(), ii)
							return
						}
					}
					message := &Message{}
					switch len(s) {
					case 1:
						message.Singular = s[0]
					case 2:
						if f.Context {
							message.Context = s[0]
							message.Singular = s[1]
						} else {
							message.Singular = s[0]
							message.Plural = s[1]
						}
					case 3:
						message.Context = s[0]
						message.Singular = s[1]
						message.Plural = s[2]
					}
					// TODO: The line number doesn't match exactly because of the
					// prepended variables
					pos := templatePosition(path, text, n)
					if err = messages.Add(message, pos, ""); err != nil {
						return
					}
				}
			}
		})
	}
	return err
}

func templatePosition(name string, text string, n parse.Node) *token.Position {
	return &token.Position{
		Filename: name,
		Line:     strings.Count(text[:int(n.Position())], "\n") + 1,
	}
}
<|endoftext|>


<commit_before>// Package help allows you to define help sections
// for your command line interfaces. It extends the flag package
// to provide flag support for help section.
//
// The package defines a subcommand "help" that you can integrate into you main CLI.
//
// For example:
//
//      $ sbr help file
//
// The subcommand displays a summary of all the sections available or the detail for any section.
//
// Sections summary example:
//
//     $ sbr help
//     Usage: sbr help

//
//     where
 is one of:
//        1.file             help relative to file management
//
// Sections details:
//     $ sbr help file dir
//     File
//
//     This document describe the file format
//     ...
//
// Sections are written using Markdown, and displayed in the terminal output using ansi escape code.
package help

import (
	"bytes"
	"flag"
	"fmt"
	"os"

	"github.com/ericaro/ansifmt/ansiblackfriday"
	"github.com/russross/blackfriday"
)

var (
	Command          = New()
	commonExtensions = 0 |
		blackfriday.EXTENSION_NO_INTRA_EMPHASIS |
		blackfriday.EXTENSION_TABLES |
		blackfriday.EXTENSION_FENCED_CODE |
		blackfriday.EXTENSION_AUTOLINK |
		blackfriday.EXTENSION_STRIKETHROUGH |
		blackfriday.EXTENSION_SPACE_HEADERS |
		blackfriday.EXTENSION_HEADER_IDS
)

//HelpCommand contains a slice of sections representing the help subsystem.
//
// HelpCommand can parse "args"
//
// and it is compatible with rakyll command package.
type HelpCommand struct {
	sections []*section
}

//New creates a new empty HelpCommand
func New() *HelpCommand {
	return &HelpCommand{
		sections: make([]*section, 0, 10),
	}
}

// unchanged flags parsing

//Flags is to update a FlagSet with the options for the helpCommand.
//
//Currently, there are no flags.
func (h *HelpCommand) Flags(fs *flag.FlagSet) *flag.FlagSet { return fs }

//Run execute the HelpCommand:
//
// with no args: prints the section summary
//
// with arguments print out every section found in the args, or an error message
func (h *HelpCommand) Run(args []string) {
	if len(args) == 0 {
		prog, helps := os.Args[0], os.Args[1]
		fmt.Fprintf(os.Stderr, "Usage: %s %s
\n\n", prog, helps)
		fmt.Fprintf(os.Stderr, "where
 is one of:\n")

		for i, s := range h.sections {
			fmt.Fprintf(os.Stderr, "  %v.%-15s %s\n", i+1, s.name, s.description)
		}
		fmt.Fprintln(os.Stderr)

	} else {
		for _, name := range args { // for each arg, that is suppose to be a "section"
			// section lookup
			exists := false
			var s *section
			for _, s = range h.sections {
				if s.name == name {
					exists = true
					break
				}
			}

			if !exists {
				fmt.Printf("No help found for %q\n", name)
			} else {
				s.Print()
			}
		}
	}
}

//Section create a new section in 'Command'.
//
// a section's name is used to identify it
//
// a section's description is used in the summary
//
// a section's content is a Markdown description.
func (h *HelpCommand) Section(name, description, content string) {
	h.sections = append(h.sections, &section{name, description, content})
}

//Section create a new section in 'Command'.
//
// a section's name is used to identify it
//
// a section's description is used in the summary
//
// a section's content is a Markdown description.
func Section(name, description, content string) {
	Command.Section(name, description, content)
}

type section struct{ name, description, content string }

//Print out the section help
func (s *section) Print() {
	renderer := ansiblackfriday.NewAnsiRenderer()

	buf := new(bytes.Buffer)
	_, err := buf.Write(blackfriday.Markdown(([]byte)(s.content), renderer, commonExtensions))
	if err != nil {
		fmt.Fprintf(os.Stderr, "buffer error: %v", err)
		return
	}
	fmt.Fprintln(os.Stderr, buf.String())
}
<commit_msg>fixed a typo in the Cmd output<commit_after>// Package help allows you to define help sections
// for your command line interfaces. It extends the flag package
// to provide flag support for help section.
//
// The package defines a subcommand "help" that you can integrate into you main CLI.
//
// For example:
//
//      $ sbr help file
//
// The subcommand displays a summary of all the sections available or the detail for any section.
//
// Sections summary example:
//
//     $ sbr help
//     Usage: sbr help

//
//     where
 is one of:
//        1.file             help relative to file management
//
// Sections details:
//     $ sbr help file dir
//     File
//
//     This document describe the file format
//     ...
//
// Sections are written using Markdown, and displayed in the terminal output using ansi escape code.
package help

import (
	"bytes"
	"flag"
	"fmt"
	"os"

	"github.com/ericaro/ansifmt/ansiblackfriday"
	"github.com/russross/blackfriday"
)

var (
	Command          = New()
	commonExtensions = 0 |
		blackfriday.EXTENSION_NO_INTRA_EMPHASIS |
		blackfriday.EXTENSION_TABLES |
		blackfriday.EXTENSION_FENCED_CODE |
		blackfriday.EXTENSION_AUTOLINK |
		blackfriday.EXTENSION_STRIKETHROUGH |
		blackfriday.EXTENSION_SPACE_HEADERS |
		blackfriday.EXTENSION_HEADER_IDS
)

//HelpCommand contains a slice of sections representing the help subsystem.
//
// HelpCommand can parse "args"
//
// and it is compatible with rakyll command package.
type HelpCommand struct {
	sections []*section
}

//New creates a new empty HelpCommand
func New() *HelpCommand {
	return &HelpCommand{
		sections: make([]*section, 0, 10),
	}
}

// unchanged flags parsing

//Flags is to update a FlagSet with the options for the helpCommand.
//
//Currently, there are no flags.
func (h *HelpCommand) Flags(fs *flag.FlagSet) *flag.FlagSet { return fs }

//Run execute the HelpCommand:
//
// with no args: prints the section summary
//
// with arguments print out every section found in the args, or an error message
func (h *HelpCommand) Run(args []string) {
	if len(args) == 0 {
		prog, helps := os.Args[0], os.Args[1]
		fmt.Fprintf(os.Stderr, "Usage: %s %s
\n\n", prog, helps)
		fmt.Fprintf(os.Stderr, "where
 is one of:\n")

		for _, s := range h.sections {
			fmt.Fprintf(os.Stderr, "  %-15s %s\n", s.name, s.description)
		}
		fmt.Fprintln(os.Stderr)

	} else {
		for _, name := range args { // for each arg, that is suppose to be a "section"
			// section lookup
			exists := false
			var s *section
			for _, s = range h.sections {
				if s.name == name {
					exists = true
					break
				}
			}

			if !exists {
				fmt.Printf("No help found for %q\n", name)
			} else {
				s.Print()
			}
		}
	}
}

//Section create a new section in 'Command'.
//
// a section's name is used to identify it
//
// a section's description is used in the summary
//
// a section's content is a Markdown description.
func (h *HelpCommand) Section(name, description, content string) {
	h.sections = append(h.sections, &section{name, description, content})
}

//Section create a new section in 'Command'.
//
// a section's name is used to identify it
//
// a section's description is used in the summary
//
// a section's content is a Markdown description.
func Section(name, description, content string) {
	Command.Section(name, description, content)
}

type section struct{ name, description, content string }

//Print out the section help
func (s *section) Print() {
	renderer := ansiblackfriday.NewAnsiRenderer()

	buf := new(bytes.Buffer)
	_, err := buf.Write(blackfriday.Markdown(([]byte)(s.content), renderer, commonExtensions))
	if err != nil {
		fmt.Fprintf(os.Stderr, "buffer error: %v", err)
		return
	}
	fmt.Fprintln(os.Stderr, buf.String())
}
<|endoftext|>


<commit_before>package slackscot

import (
	"fmt"
	"github.com/alexandre-normand/slackscot/config"
	"github.com/spf13/viper"
	"io"
	"strings"
)

type helpPlugin struct {
	Plugin

	v                      *viper.Viper
	name                   string
	slackscotVersion       string
	commands               []ActionDefinition
	hearActions            []ActionDefinition
	pluginScheduledActions []pluginScheduledAction
}

const (
	helpPluginName = "help"
)

// pluginScheduledAction represents a plugin's scheduled action with the plugin name and the action's definition
type pluginScheduledAction struct {
	plugin string
	ScheduledActionDefinition
}

func newHelpPlugin(name string, version string, viper *viper.Viper, plugins []*Plugin) *helpPlugin {
	commands, hearActions, scheduledActions := findAllActions(plugins)

	helpPlugin := new(helpPlugin)
	helpPlugin.v = viper
	helpPlugin.name = name
	helpPlugin.slackscotVersion = version
	helpPlugin.commands = commands
	helpPlugin.hearActions = hearActions
	helpPlugin.pluginScheduledActions = scheduledActions

	helpPlugin.Plugin = Plugin{Name: helpPluginName, Commands: []ActionDefinition{{
		Match: func(m *IncomingMessage) bool {
			return strings.HasPrefix(m.NormalizedText, "help")
		},
		Usage:       helpPluginName,
		Description: "Reply with usage instructions",
		Answer:      helpPlugin.showHelp,
	}}, HearActions: nil}

	return helpPlugin
}

// showHelp generates a message providing a list of all of the slackscot commands and hear actions.
// Note that ActionDefinitions with the flag Hidden set to true won't be included in the list
func (h *helpPlugin) showHelp(m *IncomingMessage) *Answer {
	var b strings.Builder

	// Get the user's first name using the botservices
	userID := m.User
	user, err := h.UserInfoFinder.GetUserInfo(userID)
	if err != nil {
		h.Logger.Debugf("Error getting user info for user id [%s] so skipping mentioning the name (it would be awkward): %v", userID, err)
	} else {
		fmt.Fprintf(&b, "🤝 You're `%s` and ", user.RealName)
	}

	fmt.Fprintf(&b, "I'm `%s` (engine `v%s`). I listen to the team's chat and provides automated functions 🧞♂.\n", h.name, h.slackscotVersion)

	if len(h.commands) > 0 {
		fmt.Fprintf(&b, "\nI currently support the following commands:\n")

		appendActions(&b, h.commands)
	}

	if len(h.hearActions) > 0 {
		fmt.Fprintf(&b, "\nAnd listen for the following:\n")

		appendActions(&b, h.hearActions)
	}

	if len(h.pluginScheduledActions) > 0 {
		fmt.Fprintf(&b, "\nAnd do those things periodically:\n")

		appendScheduledActions(&b, h.v.GetString(config.TimeLocationKey), h.pluginScheduledActions)
	}

	return &Answer{Text: b.String()}
}

func appendActions(w io.Writer, actions []ActionDefinition) {
	fmt.Fprintf(w, "\nAnd listen for the following:\n")

	for _, value := range actions {
		if value.Usage != "" && !value.Hidden {
			fmt.Fprintf(w, "\t• `%s` - %s\n", value.Usage, value.Description)
		}
	}
}

func appendScheduledActions(w io.Writer, timeLocationName string, scheduledActions []pluginScheduledAction) {
	for _, value := range scheduledActions {
		if !value.ScheduledActionDefinition.Hidden {
			fmt.Fprintf(w, "\t• [`%s`] `%s` (`%s`) - %s\n", value.plugin, value.ScheduledActionDefinition.Schedule, timeLocationName, value.ScheduledActionDefinition.Description)
		}
	}
}

func findAllActions(plugins []*Plugin) (commands []ActionDefinition, hearActions []ActionDefinition, pluginScheduledActions []pluginScheduledAction) {
	commands = make([]ActionDefinition, 0)
	hearActions = make([]ActionDefinition, 0)
	pluginScheduledActions = make([]pluginScheduledAction, 0)

	for _, p := range plugins {
		commands = append(commands, p.Commands...)
		hearActions = append(hearActions, p.HearActions...)

		if p.ScheduledActions != nil {
			for _, sa := range p.ScheduledActions {
				pluginScheduledActions = append(pluginScheduledActions, pluginScheduledAction{plugin: p.Name, ScheduledActionDefinition: sa})
			}
		}
	}

	return commands, hearActions, pluginScheduledActions
}
<commit_msg>Fix Help Message Formatting<commit_after>package slackscot

import (
	"fmt"
	"github.com/alexandre-normand/slackscot/config"
	"github.com/spf13/viper"
	"io"
	"strings"
)

type helpPlugin struct {
	Plugin

	v                      *viper.Viper
	name                   string
	slackscotVersion       string
	commands               []ActionDefinition
	hearActions            []ActionDefinition
	pluginScheduledActions []pluginScheduledAction
}

const (
	helpPluginName = "help"
)

// pluginScheduledAction represents a plugin's scheduled action with the plugin name and the action's definition
type pluginScheduledAction struct {
	plugin string
	ScheduledActionDefinition
}

func newHelpPlugin(name string, version string, viper *viper.Viper, plugins []*Plugin) *helpPlugin {
	commands, hearActions, scheduledActions := findAllActions(plugins)

	helpPlugin := new(helpPlugin)
	helpPlugin.v = viper
	helpPlugin.name = name
	helpPlugin.slackscotVersion = version
	helpPlugin.commands = commands
	helpPlugin.hearActions = hearActions
	helpPlugin.pluginScheduledActions = scheduledActions

	helpPlugin.Plugin = Plugin{Name: helpPluginName, Commands: []ActionDefinition{{
		Match: func(m *IncomingMessage) bool {
			return strings.HasPrefix(m.NormalizedText, "help")
		},
		Usage:       helpPluginName,
		Description: "Reply with usage instructions",
		Answer:      helpPlugin.showHelp,
	}}, HearActions: nil}

	return helpPlugin
}

// showHelp generates a message providing a list of all of the slackscot commands and hear actions.
// Note that ActionDefinitions with the flag Hidden set to true won't be included in the list
func (h *helpPlugin) showHelp(m *IncomingMessage) *Answer {
	var b strings.Builder

	// Get the user's first name using the botservices
	userID := m.User
	user, err := h.UserInfoFinder.GetUserInfo(userID)
	if err != nil {
		h.Logger.Debugf("Error getting user info for user id [%s] so skipping mentioning the name (it would be awkward): %v", userID, err)
	} else {
		fmt.Fprintf(&b, "🤝 You're `%s` and ", user.RealName)
	}

	fmt.Fprintf(&b, "I'm `%s` (engine `v%s`). I listen to the team's chat and provides automated functions 🧞♂.\n", h.name, h.slackscotVersion)

	if len(h.commands) > 0 {
		fmt.Fprintf(&b, "\nI currently support the following commands:\n")

		appendActions(&b, h.commands)
	}

	if len(h.hearActions) > 0 {
		fmt.Fprintf(&b, "\nAnd listen for the following:\n")

		appendActions(&b, h.hearActions)
	}

	if len(h.pluginScheduledActions) > 0 {
		fmt.Fprintf(&b, "\nAnd do those things periodically:\n")

		appendScheduledActions(&b, h.v.GetString(config.TimeLocationKey), h.pluginScheduledActions)
	}

	return &Answer{Text: b.String()}
}

func appendActions(w io.Writer, actions []ActionDefinition) {
	for _, value := range actions {
		if value.Usage != "" && !value.Hidden {
			fmt.Fprintf(w, "\t• `%s` - %s\n", value.Usage, value.Description)
		}
	}
}

func appendScheduledActions(w io.Writer, timeLocationName string, scheduledActions []pluginScheduledAction) {
	for _, value := range scheduledActions {
		if !value.ScheduledActionDefinition.Hidden {
			fmt.Fprintf(w, "\t• [`%s`] `%s` (`%s`) - %s\n", value.plugin, value.ScheduledActionDefinition.Schedule, timeLocationName, value.ScheduledActionDefinition.Description)
		}
	}
}

func findAllActions(plugins []*Plugin) (commands []ActionDefinition, hearActions []ActionDefinition, pluginScheduledActions []pluginScheduledAction) {
	commands = make([]ActionDefinition, 0)
	hearActions = make([]ActionDefinition, 0)
	pluginScheduledActions = make([]pluginScheduledAction, 0)

	for _, p := range plugins {
		commands = append(commands, p.Commands...)
		hearActions = append(hearActions, p.HearActions...)

		if p.ScheduledActions != nil {
			for _, sa := range p.ScheduledActions {
				pluginScheduledActions = append(pluginScheduledActions, pluginScheduledAction{plugin: p.Name, ScheduledActionDefinition: sa})
			}
		}
	}

	return commands, hearActions, pluginScheduledActions
}
<|endoftext|>


<commit_before>package gonfig

import (
	"fmt"
	"io"
	"os"
	"path"
	"reflect"
	"strings"
)

const (
	defaultHelpDescription = "print this help menu"
	defaultHelpMessage     = "Usage of __EXEC__:"
)

func typeString(t reflect.Type) string {
	if t.Implements(typeOfTextUnmarshaler) {
		return "string"
	}

	if t == typeOfByteSlice {
		return "string"
	}

	switch t.Kind() {
	case reflect.String:
		return "string"

	case reflect.Bool:
		return "bool"

	case reflect.Float32, reflect.Float64:
		return "float"

	case reflect.Int, reflect.Int64, reflect.Int32, reflect.Int16, reflect.Int8:
		return "int"

	case reflect.Uint, reflect.Uint64, reflect.Uint32, reflect.Uint16, reflect.Uint8:
		return "uint"

	case reflect.Slice:
		subTypeStr := typeString(t.Elem())
		return subTypeStr + "..."
	}

	return ""
}

// unquoteDescription extracts a back-quoted name from the description
// string and returns it and the un-quoted description.
// Given "a `name` to show" it returns ("name", "a name to show").
// If there are no back quotes, the name is empty.
func unquoteDescription(desc string) (string, string) {
	for i := 0; i < len(desc); i++ {
		if desc[i] == '`' {
			for j := i + 1; j < len(desc); j++ {
				if desc[j] == '`' {
					name := desc[i+1 : j]
					newDesc := desc[:i] + name + desc[j+1:]
					return name, newDesc
				}
			}
			break // only one back quote
		}
	}

	return "", desc
}

// writeHelpMessage writes the help message to the writer.
//
// This implementation is borrowed from https://github.com/spf13/pflag
func writeHelpMessage(s *setup, w io.Writer) {
	lines := make([]string, 0, len(s.allOpts))

	maxlen := 0
	for _, opt := range s.allOpts {
		if opt.isParent {
			continue
		}

		line := ""
		if opt.short != "" {
			line = fmt.Sprintf("  -%s, --%s", opt.short, opt.fullID())
		} else {
			line = fmt.Sprintf("      --%s", opt.fullID())
		}

		typeStr := typeString(opt.value.Type())
		varname, desc := unquoteDescription(opt.desc)
		if varname == "" {
			varname = typeStr
			if varname == "bool" {
				// We don't want to show a varname for bools.
				varname = ""
			}
		}

		if varname != "" {
			line += " " + varname
		}

		// This special character will be replaced with spacing once the
		// correct alignment is calculated
		line += "\x00"
		if len(line) > maxlen {
			maxlen = len(line)
		}

		line += desc
		if opt.defaul != "" {
			if len(typeStr) >= 6 && typeStr[0:6] == "string" {
				line += fmt.Sprintf(" (default %q)", opt.defaul)
			} else {
				line += fmt.Sprintf(" (default %s)", opt.defaul)
			}
		}

		lines = append(lines, line)
	}

	helpFlagDesc := s.conf.HelpDescription
	if helpFlagDesc == "" {
		helpFlagDesc = defaultHelpDescription
	}
	lines = append(lines, "  -h, --help\x00"+helpFlagDesc)

	message := s.conf.HelpMessage
	if message == "" {
		exec := path.Base(os.Args[0])
		message = strings.Replace(defaultHelpMessage, "__EXEC__", exec, 1)
	}
	fmt.Fprintln(w, message)

	for _, line := range lines {
		sidx := strings.Index(line, "\x00")
		spacing := strings.Repeat(" ", maxlen-sidx)
		// maxlen + 2 comes from + 1 for the \x00 and + 1 for the (deliberate)
		// off-by-one in maxlen-sidx
		fmt.Fprintln(w, line[:sidx], spacing, line[sidx+1:])
	}
}

// printHelpAndExit prints the help message and exits the program.
func printHelpAndExit(s *setup) {
	writeHelpMessage(s, os.Stdout)
	os.Exit(2)
}
<commit_msg>Better help formatting for maps<commit_after>package gonfig

import (
	"fmt"
	"io"
	"os"
	"path"
	"reflect"
	"strings"
)

const (
	defaultHelpDescription = "print this help menu"
	defaultHelpMessage     = "Usage of __EXEC__:"
)

func typeString(t reflect.Type) string {
	if t.Implements(typeOfTextUnmarshaler) {
		return "string"
	}

	if t == typeOfByteSlice {
		return "string"
	}

	switch t.Kind() {
	case reflect.String:
		return "string"

	case reflect.Bool:
		return "bool"

	case reflect.Float32, reflect.Float64:
		return "float"

	case reflect.Int, reflect.Int64, reflect.Int32, reflect.Int16, reflect.Int8:
		return "int"

	case reflect.Uint, reflect.Uint64, reflect.Uint32, reflect.Uint16, reflect.Uint8:
		return "uint"

	case reflect.Slice:
		subTypeStr := typeString(t.Elem())
		return subTypeStr + "..."
	}

	return ""
}

// unquoteDescription extracts a back-quoted name from the description
// string and returns it and the un-quoted description.
// Given "a `name` to show" it returns ("name", "a name to show").
// If there are no back quotes, the name is empty.
func unquoteDescription(desc string) (string, string) {
	for i := 0; i < len(desc); i++ {
		if desc[i] == '`' {
			for j := i + 1; j < len(desc); j++ {
				if desc[j] == '`' {
					name := desc[i+1 : j]
					newDesc := desc[:i] + name + desc[j+1:]
					return name, newDesc
				}
			}
			break // only one back quote
		}
	}

	return "", desc
}

// writeHelpMessage writes the help message to the writer.
//
// This implementation is borrowed from https://github.com/spf13/pflag
func writeHelpMessage(s *setup, w io.Writer) {
	lines := make([]string, 0, len(s.allOpts))

	maxlen := 0
	for _, opt := range s.allOpts {
		if opt.isParent {
			continue
		}

		line := ""
		if opt.short != "" {
			line = fmt.Sprintf("  -%s, --%s", opt.short, opt.fullID())
		} else {
			line = fmt.Sprintf("      --%s", opt.fullID())
		}

		typeStr := typeString(opt.value.Type())
		varname, desc := unquoteDescription(opt.desc)
		if opt.isMap {
			line += ".<key> <value>"
		} else {
			if varname == "" {
				varname = typeStr
				if varname == "bool" {
					// We don't want to show a varname for bools.
					varname = ""
				}
			}

			if varname != "" {
				line += " " + varname
			}
		}

		// This special character will be replaced with spacing once the
		// correct alignment is calculated
		line += "\x00"
		if len(line) > maxlen {
			maxlen = len(line)
		}

		line += desc
		if opt.defaul != "" {
			if len(typeStr) >= 6 && typeStr[0:6] == "string" {
				// Put quotes around string types.
				line += fmt.Sprintf(" (default %q)", opt.defaul)
			} else {
				line += fmt.Sprintf(" (default %s)", opt.defaul)
			}
		}

		lines = append(lines, line)
	}

	helpFlagDesc := s.conf.HelpDescription
	if helpFlagDesc == "" {
		helpFlagDesc = defaultHelpDescription
	}
	lines = append(lines, "  -h, --help\x00"+helpFlagDesc)

	message := s.conf.HelpMessage
	if message == "" {
		exec := path.Base(os.Args[0])
		message = strings.Replace(defaultHelpMessage, "__EXEC__", exec, 1)
	}
	fmt.Fprintln(w, message)

	for _, line := range lines {
		sidx := strings.Index(line, "\x00")
		spacing := strings.Repeat(" ", maxlen-sidx)
		// maxlen + 2 comes from + 1 for the \x00 and + 1 for the (deliberate)
		// off-by-one in maxlen-sidx
		fmt.Fprintln(w, line[:sidx], spacing, line[sidx+1:])
	}
}

// printHelpAndExit prints the help message and exits the program.
func printHelpAndExit(s *setup) {
	writeHelpMessage(s, os.Stdout)
	os.Exit(2)
}
<|endoftext|>