<commit_before>package make

import (
	log "github.com/Sirupsen/logrus"
	"os"
	"os/exec"
	"strings"
)

type SolrCore struct {
	Address    string
	Name       string
	Template   string
	Path       string
	DataPath   string
	ConfigFile string
	SchemaFile string
}

func logSolrInstall() bool {
	if verifySolrInstall() {
		log.Infoln("Found Solr installation")
		return true
	} else {
		log.Errorln("Could not find Solr installation")
		return false
	}
}
func logSolrCLI() bool {
	if verifySolrCLI() {
		log.Infoln("Found Solr command-line tools")
		return true
	} else {
		log.Warnln("Could not find Solr command-line tools")
		return false
	}
}
func logResources(Template string) bool {
	if verifyResources(Template) {
		log.Infoln("Found configuration folder")
		return true
	} else {
		log.Errorln("Could not find configuration folder")
		return false
	}
}
func logSolrCore(SolrCore *SolrCore) bool {
	if verifySolrCore(SolrCore) {
		log.Infoln("Solr core is installed.")
		return true
	} else {
		log.Warnln("Solr core is not installed.")
		return false
	}
}
func verifySolrInstall() bool {
	_, err := os.Stat("/opt/solr")
	if err == nil {
		return true
	} else {
		return false
	}
}
func verifySolrCLI() bool {
	_, err := os.Stat("/opt/solr/bin/solr")
	if err == nil {
		return true
	} else {
		return false
	}
}
func verifyResources(Template string) bool {
	_, err := os.Stat(Template)
	if err == nil {
		return true
	} else {
		return false
	}
}
func verifySolrCore(SolrCore *SolrCore) bool {
	curlResponse, err := exec.Command("curl", SolrCore.Address+"/solr/admin/cores?action=STATUS").Output()
	if err == nil {
		if strings.Contains(string(curlResponse), `<str name="name">`+SolrCore.Name+`</str>`) == true {
			return true
		} else {
			return false
		}
	} else {
		log.Errorln("Solr could not be accessed using CURL:", err.Error())
	}
	return false
}

func NewCore(Address, Name, Template, Path, DataPath, ConfigFile, SchemaFile string) SolrCore {
	return SolrCore{Address, Name, Template, Path, DataPath, ConfigFile, SchemaFile}
}

func (SolrCore *SolrCore) Install() {
	if logSolrInstall() && logResources(SolrCore.Template) {
		log.Infoln("All checks have passed.")
		dataDir := SolrCore.Path + "/" + SolrCore.DataPath + "/" + SolrCore.Name + "/conf/"

		// Create data directories
		err := os.MkdirAll(dataDir, 0777)
		if err == nil {
			log.Infoln("Directory has been created.", dataDir)
		} else {
			log.Errorln("Directory has not been created:", err.Error())
		}

		// Sync
		_, err = exec.Command("rsync", "-a", SolrCore.Template+"/", dataDir).Output()
		if err == nil {
			log.Infoln("Configuration has been synced with boilerplate resources.")
		} else {
			log.Errorln("Configuration could not be synced with boilerplate resources:", err.Error())
		}

		_, err = exec.Command("curl", SolrCore.Address+"/solr/admin/cores?action=CREATE&name="+SolrCore.Name+"&instanceDir="+SolrCore.Name+"&dataDir="+SolrCore.DataPath+"&config="+SolrCore.ConfigFile+"&schema="+SolrCore.SchemaFile).Output()
		if err == nil {
			log.Infoln("Core has been successfully installed.")
		} else {
			log.Errorln("Core could not be installed:", err)
		}
	}
	verifySolrCore(SolrCore)
}

func (SolrCore *SolrCore) Uninstall() {

	dataDir := SolrCore.Path + "/" + SolrCore.DataPath + "/" + SolrCore.Name

	_, err := exec.Command("curl", SolrCore.Address+"/solr/admin/cores?action=UNLOAD&core="+SolrCore.Name).Output()
	if err == nil {
		log.Infoln("Core has been successfully uninstalled.")
	} else {
		log.Errorln("Core could not be uninstalled:", err)
	}
	err = os.RemoveAll(dataDir)
	if err == nil {
		log.Infoln("Core resources have been removed.")
	} else {
		log.Errorln("Core resources could not be removed:", err)
	}
}
<commit_msg>Create a data folder when conf folder is created.<commit_after>package make

import (
	log "github.com/Sirupsen/logrus"
	"os"
	"os/exec"
	"strings"
)

type SolrCore struct {
	Address    string
	Name       string
	Template   string
	Path       string
	DataPath   string
	ConfigFile string
	SchemaFile string
}

func logSolrInstall() bool {
	if verifySolrInstall() {
		log.Infoln("Found Solr installation")
		return true
	} else {
		log.Errorln("Could not find Solr installation")
		return false
	}
}
func logSolrCLI() bool {
	if verifySolrCLI() {
		log.Infoln("Found Solr command-line tools")
		return true
	} else {
		log.Warnln("Could not find Solr command-line tools")
		return false
	}
}
func logResources(Template string) bool {
	if verifyResources(Template) {
		log.Infoln("Found configuration folder")
		return true
	} else {
		log.Errorln("Could not find configuration folder")
		return false
	}
}
func logSolrCore(SolrCore *SolrCore) bool {
	if verifySolrCore(SolrCore) {
		log.Infoln("Solr core is installed.")
		return true
	} else {
		log.Warnln("Solr core is not installed.")
		return false
	}
}
func verifySolrInstall() bool {
	_, err := os.Stat("/opt/solr")
	if err == nil {
		return true
	} else {
		return false
	}
}
func verifySolrCLI() bool {
	_, err := os.Stat("/opt/solr/bin/solr")
	if err == nil {
		return true
	} else {
		return false
	}
}
func verifyResources(Template string) bool {
	_, err := os.Stat(Template)
	if err == nil {
		return true
	} else {
		return false
	}
}
func verifySolrCore(SolrCore *SolrCore) bool {
	curlResponse, err := exec.Command("curl", SolrCore.Address+"/solr/admin/cores?action=STATUS").Output()
	if err == nil {
		if strings.Contains(string(curlResponse), `<str name="name">`+SolrCore.Name+`</str>`) == true {
			return true
		} else {
			return false
		}
	} else {
		log.Errorln("Solr could not be accessed using CURL:", err.Error())
	}
	return false
}

func NewCore(Address, Name, Template, Path, DataPath, ConfigFile, SchemaFile string) SolrCore {
	return SolrCore{Address, Name, Template, Path, DataPath, ConfigFile, SchemaFile}
}

func (SolrCore *SolrCore) Install() {
	if logSolrInstall() && logResources(SolrCore.Template) {
		log.Infoln("All checks have passed.")
		dataDir := SolrCore.Path + "/" + SolrCore.DataPath + "/" + SolrCore.Name + "/conf/"

		// Create data directories
		err := os.MkdirAll(SolrCore.Path+"/"+SolrCore.DataPath+"/"+SolrCore.Name+"/data/", 0777)
		if err == nil {
			log.Infoln("Directory has been created.", SolrCore.Path+"/"+SolrCore.DataPath+"/"+SolrCore.Name+"/data/")
		} else {
			log.Errorln("Directory has not been created:", err.Error())
		}

		err = os.MkdirAll(dataDir, 0777)
		if err == nil {
			log.Infoln("Directory has been created.", dataDir)
		} else {
			log.Errorln("Directory has not been created:", err.Error())
		}

		// Sync
		_, err = exec.Command("rsync", "-a", SolrCore.Template+"/", dataDir).Output()
		if err == nil {
			log.Infoln("Configuration has been synced with boilerplate resources.")
		} else {
			log.Errorln("Configuration could not be synced with boilerplate resources:", err.Error())
		}

		_, err = exec.Command("curl", SolrCore.Address+"/solr/admin/cores?action=CREATE&name="+SolrCore.Name+"&instanceDir="+SolrCore.Name+"&dataDir="+SolrCore.DataPath+"&config="+SolrCore.ConfigFile+"&schema="+SolrCore.SchemaFile).Output()
		if err == nil {
			log.Infoln("Core has been successfully installed.")
		} else {
			log.Errorln("Core could not be installed:", err)
		}
	}
	verifySolrCore(SolrCore)
}

func (SolrCore *SolrCore) Uninstall() {

	dataDir := SolrCore.Path + "/" + SolrCore.DataPath + "/" + SolrCore.Name

	_, err := exec.Command("curl", SolrCore.Address+"/solr/admin/cores?action=UNLOAD&core="+SolrCore.Name).Output()
	if err == nil {
		log.Infoln("Core has been successfully uninstalled.")
	} else {
		log.Errorln("Core could not be uninstalled:", err)
	}
	err = os.RemoveAll(dataDir)
	if err == nil {
		log.Infoln("Core resources have been removed.")
	} else {
		log.Errorln("Core resources could not be removed:", err)
	}
}
<|endoftext|>


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

/*
Package net provides a portable interface for network I/O, including
TCP/IP, UDP, domain name resolution, and Unix domain sockets.

Although the package provides access to low-level networking
primitives, most clients will need only the basic interface provided
by the Dial, Listen, and Accept functions and the associated
Conn and Listener interfaces. The crypto/tls package uses
the same interfaces and similar Dial and Listen functions.

The Dial function connects to a server:

	conn, err := net.Dial("tcp", "google.com:80")
	if err != nil {
		// handle error
	}
	fmt.Fprintf(conn, "GET / HTTP/1.0\r\n\r\n")
	status, err := bufio.NewReader(conn).ReadString('\n')
	// ...

The Listen function creates servers:

	ln, err := net.Listen("tcp", ":8080")
	if err != nil {
		// handle error
	}
	for {
		conn, err := ln.Accept()
		if err != nil {
			// handle error
			continue
		}
		go handleConnection(conn)
	}
*/
package net

// TODO(rsc):
//	support for raw ethernet sockets

import (
	"errors"
	"io"
	"os"
	"syscall"
	"time"
)

// Addr represents a network end point address.
type Addr interface {
	Network() string // name of the network
	String() string  // string form of address
}

// Conn is a generic stream-oriented network connection.
//
// Multiple goroutines may invoke methods on a Conn simultaneously.
type Conn interface {
	// Read reads data from the connection.
	// Read can be made to time out and return a Error with Timeout() == true
	// after a fixed time limit; see SetDeadline and SetReadDeadline.
	Read(b []byte) (n int, err error)

	// Write writes data to the connection.
	// Write can be made to time out and return a Error with Timeout() == true
	// after a fixed time limit; see SetDeadline and SetWriteDeadline.
	Write(b []byte) (n int, err error)

	// Close closes the connection.
	// Any blocked Read or Write operations will be unblocked and return errors.
	Close() error

	// LocalAddr returns the local network address.
	LocalAddr() Addr

	// RemoteAddr returns the remote network address.
	RemoteAddr() Addr

	// SetDeadline sets the read and write deadlines associated
	// with the connection. It is equivalent to calling both
	// SetReadDeadline and SetWriteDeadline.
	//
	// A deadline is an absolute time after which I/O operations
	// fail with a timeout (see type Error) instead of
	// blocking. The deadline applies to all future I/O, not just
	// the immediately following call to Read or Write.
	//
	// An idle timeout can be implemented by repeatedly extending
	// the deadline after successful Read or Write calls.
	//
	// A zero value for t means I/O operations will not time out.
	SetDeadline(t time.Time) error

	// SetReadDeadline sets the deadline for future Read calls.
	// A zero value for t means Read will not time out.
	SetReadDeadline(t time.Time) error

	// SetWriteDeadline sets the deadline for future Write calls.
	// Even if write times out, it may return n > 0, indicating that
	// some of the data was successfully written.
	// A zero value for t means Write will not time out.
	SetWriteDeadline(t time.Time) error
}

type conn struct {
	fd *netFD
}

func (c *conn) ok() bool { return c != nil && c.fd != nil }

// Implementation of the Conn interface.

// Read implements the Conn Read method.
func (c *conn) Read(b []byte) (int, error) {
	if !c.ok() {
		return 0, syscall.EINVAL
	}
	return c.fd.Read(b)
}

// Write implements the Conn Write method.
func (c *conn) Write(b []byte) (int, error) {
	if !c.ok() {
		return 0, syscall.EINVAL
	}
	return c.fd.Write(b)
}

// Close closes the connection.
func (c *conn) Close() error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return c.fd.Close()
}

// LocalAddr returns the local network address.
func (c *conn) LocalAddr() Addr {
	if !c.ok() {
		return nil
	}
	return c.fd.laddr
}

// RemoteAddr returns the remote network address.
func (c *conn) RemoteAddr() Addr {
	if !c.ok() {
		return nil
	}
	return c.fd.raddr
}

// SetDeadline implements the Conn SetDeadline method.
func (c *conn) SetDeadline(t time.Time) error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return c.fd.setDeadline(t)
}

// SetReadDeadline implements the Conn SetReadDeadline method.
func (c *conn) SetReadDeadline(t time.Time) error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return c.fd.setReadDeadline(t)
}

// SetWriteDeadline implements the Conn SetWriteDeadline method.
func (c *conn) SetWriteDeadline(t time.Time) error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return c.fd.setWriteDeadline(t)
}

// SetReadBuffer sets the size of the operating system's
// receive buffer associated with the connection.
func (c *conn) SetReadBuffer(bytes int) error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return setReadBuffer(c.fd, bytes)
}

// SetWriteBuffer sets the size of the operating system's
// transmit buffer associated with the connection.
func (c *conn) SetWriteBuffer(bytes int) error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return setWriteBuffer(c.fd, bytes)
}

// File sets the underlying os.File to blocking mode and returns a copy.
// It is the caller's responsibility to close f when finished.
// Closing c does not affect f, and closing f does not affect c.
//
// The returned os.File's file descriptor is different from the connection's.
// Attempting to change properties of the original using this duplicate
// may or may not have the desired effect.
func (c *conn) File() (f *os.File, err error) { return c.fd.dup() }

// An Error represents a network error.
type Error interface {
	error
	Timeout() bool   // Is the error a timeout?
	Temporary() bool // Is the error temporary?
}

// PacketConn is a generic packet-oriented network connection.
//
// Multiple goroutines may invoke methods on a PacketConn simultaneously.
type PacketConn interface {
	// ReadFrom reads a packet from the connection,
	// copying the payload into b.  It returns the number of
	// bytes copied into b and the return address that
	// was on the packet.
	// ReadFrom can be made to time out and return
	// an error with Timeout() == true after a fixed time limit;
	// see SetDeadline and SetReadDeadline.
	ReadFrom(b []byte) (n int, addr Addr, err error)

	// WriteTo writes a packet with payload b to addr.
	// WriteTo can be made to time out and return
	// an error with Timeout() == true after a fixed time limit;
	// see SetDeadline and SetWriteDeadline.
	// On packet-oriented connections, write timeouts are rare.
	WriteTo(b []byte, addr Addr) (n int, err error)

	// Close closes the connection.
	// Any blocked ReadFrom or WriteTo operations will be unblocked and return errors.
	Close() error

	// LocalAddr returns the local network address.
	LocalAddr() Addr

	// SetDeadline sets the read and write deadlines associated
	// with the connection.
	SetDeadline(t time.Time) error

	// SetReadDeadline sets the deadline for future Read calls.
	// If the deadline is reached, Read will fail with a timeout
	// (see type Error) instead of blocking.
	// A zero value for t means Read will not time out.
	SetReadDeadline(t time.Time) error

	// SetWriteDeadline sets the deadline for future Write calls.
	// If the deadline is reached, Write will fail with a timeout
	// (see type Error) instead of blocking.
	// A zero value for t means Write will not time out.
	// Even if write times out, it may return n > 0, indicating that
	// some of the data was successfully written.
	SetWriteDeadline(t time.Time) error
}

var listenerBacklog = maxListenerBacklog()

// A Listener is a generic network listener for stream-oriented protocols.
//
// Multiple goroutines may invoke methods on a Listener simultaneously.
type Listener interface {
	// Accept waits for and returns the next connection to the listener.
	Accept() (c Conn, err error)

	// Close closes the listener.
	// Any blocked Accept operations will be unblocked and return errors.
	Close() error

	// Addr returns the listener's network address.
	Addr() Addr
}

// Various errors contained in OpError.
var (
	// For connection setup and write operations.
	errMissingAddress = errors.New("missing address")

	// For both read and write operations.
	errTimeout          error = &timeoutError{}
	errClosing                = errors.New("use of closed network connection")
	ErrWriteToConnected       = errors.New("use of WriteTo with pre-connected connection")
)

// OpError is the error type usually returned by functions in the net
// package. It describes the operation, network type, and address of
// an error.
type OpError struct {
	// Op is the operation which caused the error, such as
	// "read" or "write".
	Op string

	// Net is the network type on which this error occurred,
	// such as "tcp" or "udp6".
	Net string

	// Addr is the network address on which this error occurred.
	Addr Addr

	// Err is the error that occurred during the operation.
	Err error
}

func (e *OpError) Error() string {
	if e == nil {
		return "<nil>"
	}
	s := e.Op
	if e.Net != "" {
		s += " " + e.Net
	}
	if e.Addr != nil {
		s += " " + e.Addr.String()
	}
	s += ": " + e.Err.Error()
	return s
}

type temporary interface {
	Temporary() bool
}

func (e *OpError) Temporary() bool {
	t, ok := e.Err.(temporary)
	return ok && t.Temporary()
}

var noDeadline = time.Time{}

type timeout interface {
	Timeout() bool
}

func (e *OpError) Timeout() bool {
	t, ok := e.Err.(timeout)
	return ok && t.Timeout()
}

type timeoutError struct{}

func (e *timeoutError) Error() string   { return "i/o timeout" }
func (e *timeoutError) Timeout() bool   { return true }
func (e *timeoutError) Temporary() bool { return true }

type AddrError struct {
	Err  string
	Addr string
}

func (e *AddrError) Error() string {
	if e == nil {
		return "<nil>"
	}
	s := e.Err
	if e.Addr != "" {
		s += " " + e.Addr
	}
	return s
}

func (e *AddrError) Temporary() bool {
	return false
}

func (e *AddrError) Timeout() bool {
	return false
}

type UnknownNetworkError string

func (e UnknownNetworkError) Error() string   { return "unknown network " + string(e) }
func (e UnknownNetworkError) Temporary() bool { return false }
func (e UnknownNetworkError) Timeout() bool   { return false }

type InvalidAddrError string

func (e InvalidAddrError) Error() string   { return string(e) }
func (e InvalidAddrError) Timeout() bool   { return false }
func (e InvalidAddrError) Temporary() bool { return false }

// DNSConfigError represents an error reading the machine's DNS configuration.
type DNSConfigError struct {
	Err error
}

func (e *DNSConfigError) Error() string {
	return "error reading DNS config: " + e.Err.Error()
}

func (e *DNSConfigError) Timeout() bool   { return false }
func (e *DNSConfigError) Temporary() bool { return false }

type writerOnly struct {
	io.Writer
}

// Fallback implementation of io.ReaderFrom's ReadFrom, when sendfile isn't
// applicable.
func genericReadFrom(w io.Writer, r io.Reader) (n int64, err error) {
	// Use wrapper to hide existing r.ReadFrom from io.Copy.
	return io.Copy(writerOnly{w}, r)
}

// Limit the number of concurrent cgo-using goroutines, because
// each will block an entire operating system thread. The usual culprit
// is resolving many DNS names in separate goroutines but the DNS
// server is not responding. Then the many lookups each use a different
// thread, and the system or the program runs out of threads.

var threadLimit = make(chan struct{}, 500)

// Using send for acquire is fine here because we are not using this
// to protect any memory. All we care about is the number of goroutines
// making calls at a time.

func acquireThread() {
	threadLimit <- struct{}{}
}

func releaseThread() {
	<-threadLimit
}
<commit_msg>net: fix misleading package comment example<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
Package net provides a portable interface for network I/O, including
TCP/IP, UDP, domain name resolution, and Unix domain sockets.

Although the package provides access to low-level networking
primitives, most clients will need only the basic interface provided
by the Dial, Listen, and Accept functions and the associated
Conn and Listener interfaces. The crypto/tls package uses
the same interfaces and similar Dial and Listen functions.

The Dial function connects to a server:

	conn, err := net.Dial("tcp", "google.com:80")
	if err != nil {
		// handle error
	}
	fmt.Fprintf(conn, "GET / HTTP/1.0\r\n\r\n")
	status, err := bufio.NewReader(conn).ReadString('\n')
	// ...

The Listen function creates servers:

	ln, err := net.Listen("tcp", ":8080")
	if err != nil {
		// handle error
	}
	for {
		conn, err := ln.Accept()
		if err != nil {
			// handle error
		}
		go handleConnection(conn)
	}
*/
package net

// TODO(rsc):
//	support for raw ethernet sockets

import (
	"errors"
	"io"
	"os"
	"syscall"
	"time"
)

// Addr represents a network end point address.
type Addr interface {
	Network() string // name of the network
	String() string  // string form of address
}

// Conn is a generic stream-oriented network connection.
//
// Multiple goroutines may invoke methods on a Conn simultaneously.
type Conn interface {
	// Read reads data from the connection.
	// Read can be made to time out and return a Error with Timeout() == true
	// after a fixed time limit; see SetDeadline and SetReadDeadline.
	Read(b []byte) (n int, err error)

	// Write writes data to the connection.
	// Write can be made to time out and return a Error with Timeout() == true
	// after a fixed time limit; see SetDeadline and SetWriteDeadline.
	Write(b []byte) (n int, err error)

	// Close closes the connection.
	// Any blocked Read or Write operations will be unblocked and return errors.
	Close() error

	// LocalAddr returns the local network address.
	LocalAddr() Addr

	// RemoteAddr returns the remote network address.
	RemoteAddr() Addr

	// SetDeadline sets the read and write deadlines associated
	// with the connection. It is equivalent to calling both
	// SetReadDeadline and SetWriteDeadline.
	//
	// A deadline is an absolute time after which I/O operations
	// fail with a timeout (see type Error) instead of
	// blocking. The deadline applies to all future I/O, not just
	// the immediately following call to Read or Write.
	//
	// An idle timeout can be implemented by repeatedly extending
	// the deadline after successful Read or Write calls.
	//
	// A zero value for t means I/O operations will not time out.
	SetDeadline(t time.Time) error

	// SetReadDeadline sets the deadline for future Read calls.
	// A zero value for t means Read will not time out.
	SetReadDeadline(t time.Time) error

	// SetWriteDeadline sets the deadline for future Write calls.
	// Even if write times out, it may return n > 0, indicating that
	// some of the data was successfully written.
	// A zero value for t means Write will not time out.
	SetWriteDeadline(t time.Time) error
}

type conn struct {
	fd *netFD
}

func (c *conn) ok() bool { return c != nil && c.fd != nil }

// Implementation of the Conn interface.

// Read implements the Conn Read method.
func (c *conn) Read(b []byte) (int, error) {
	if !c.ok() {
		return 0, syscall.EINVAL
	}
	return c.fd.Read(b)
}

// Write implements the Conn Write method.
func (c *conn) Write(b []byte) (int, error) {
	if !c.ok() {
		return 0, syscall.EINVAL
	}
	return c.fd.Write(b)
}

// Close closes the connection.
func (c *conn) Close() error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return c.fd.Close()
}

// LocalAddr returns the local network address.
func (c *conn) LocalAddr() Addr {
	if !c.ok() {
		return nil
	}
	return c.fd.laddr
}

// RemoteAddr returns the remote network address.
func (c *conn) RemoteAddr() Addr {
	if !c.ok() {
		return nil
	}
	return c.fd.raddr
}

// SetDeadline implements the Conn SetDeadline method.
func (c *conn) SetDeadline(t time.Time) error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return c.fd.setDeadline(t)
}

// SetReadDeadline implements the Conn SetReadDeadline method.
func (c *conn) SetReadDeadline(t time.Time) error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return c.fd.setReadDeadline(t)
}

// SetWriteDeadline implements the Conn SetWriteDeadline method.
func (c *conn) SetWriteDeadline(t time.Time) error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return c.fd.setWriteDeadline(t)
}

// SetReadBuffer sets the size of the operating system's
// receive buffer associated with the connection.
func (c *conn) SetReadBuffer(bytes int) error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return setReadBuffer(c.fd, bytes)
}

// SetWriteBuffer sets the size of the operating system's
// transmit buffer associated with the connection.
func (c *conn) SetWriteBuffer(bytes int) error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return setWriteBuffer(c.fd, bytes)
}

// File sets the underlying os.File to blocking mode and returns a copy.
// It is the caller's responsibility to close f when finished.
// Closing c does not affect f, and closing f does not affect c.
//
// The returned os.File's file descriptor is different from the connection's.
// Attempting to change properties of the original using this duplicate
// may or may not have the desired effect.
func (c *conn) File() (f *os.File, err error) { return c.fd.dup() }

// An Error represents a network error.
type Error interface {
	error
	Timeout() bool   // Is the error a timeout?
	Temporary() bool // Is the error temporary?
}

// PacketConn is a generic packet-oriented network connection.
//
// Multiple goroutines may invoke methods on a PacketConn simultaneously.
type PacketConn interface {
	// ReadFrom reads a packet from the connection,
	// copying the payload into b.  It returns the number of
	// bytes copied into b and the return address that
	// was on the packet.
	// ReadFrom can be made to time out and return
	// an error with Timeout() == true after a fixed time limit;
	// see SetDeadline and SetReadDeadline.
	ReadFrom(b []byte) (n int, addr Addr, err error)

	// WriteTo writes a packet with payload b to addr.
	// WriteTo can be made to time out and return
	// an error with Timeout() == true after a fixed time limit;
	// see SetDeadline and SetWriteDeadline.
	// On packet-oriented connections, write timeouts are rare.
	WriteTo(b []byte, addr Addr) (n int, err error)

	// Close closes the connection.
	// Any blocked ReadFrom or WriteTo operations will be unblocked and return errors.
	Close() error

	// LocalAddr returns the local network address.
	LocalAddr() Addr

	// SetDeadline sets the read and write deadlines associated
	// with the connection.
	SetDeadline(t time.Time) error

	// SetReadDeadline sets the deadline for future Read calls.
	// If the deadline is reached, Read will fail with a timeout
	// (see type Error) instead of blocking.
	// A zero value for t means Read will not time out.
	SetReadDeadline(t time.Time) error

	// SetWriteDeadline sets the deadline for future Write calls.
	// If the deadline is reached, Write will fail with a timeout
	// (see type Error) instead of blocking.
	// A zero value for t means Write will not time out.
	// Even if write times out, it may return n > 0, indicating that
	// some of the data was successfully written.
	SetWriteDeadline(t time.Time) error
}

var listenerBacklog = maxListenerBacklog()

// A Listener is a generic network listener for stream-oriented protocols.
//
// Multiple goroutines may invoke methods on a Listener simultaneously.
type Listener interface {
	// Accept waits for and returns the next connection to the listener.
	Accept() (c Conn, err error)

	// Close closes the listener.
	// Any blocked Accept operations will be unblocked and return errors.
	Close() error

	// Addr returns the listener's network address.
	Addr() Addr
}

// Various errors contained in OpError.
var (
	// For connection setup and write operations.
	errMissingAddress = errors.New("missing address")

	// For both read and write operations.
	errTimeout          error = &timeoutError{}
	errClosing                = errors.New("use of closed network connection")
	ErrWriteToConnected       = errors.New("use of WriteTo with pre-connected connection")
)

// OpError is the error type usually returned by functions in the net
// package. It describes the operation, network type, and address of
// an error.
type OpError struct {
	// Op is the operation which caused the error, such as
	// "read" or "write".
	Op string

	// Net is the network type on which this error occurred,
	// such as "tcp" or "udp6".
	Net string

	// Addr is the network address on which this error occurred.
	Addr Addr

	// Err is the error that occurred during the operation.
	Err error
}

func (e *OpError) Error() string {
	if e == nil {
		return "<nil>"
	}
	s := e.Op
	if e.Net != "" {
		s += " " + e.Net
	}
	if e.Addr != nil {
		s += " " + e.Addr.String()
	}
	s += ": " + e.Err.Error()
	return s
}

type temporary interface {
	Temporary() bool
}

func (e *OpError) Temporary() bool {
	t, ok := e.Err.(temporary)
	return ok && t.Temporary()
}

var noDeadline = time.Time{}

type timeout interface {
	Timeout() bool
}

func (e *OpError) Timeout() bool {
	t, ok := e.Err.(timeout)
	return ok && t.Timeout()
}

type timeoutError struct{}

func (e *timeoutError) Error() string   { return "i/o timeout" }
func (e *timeoutError) Timeout() bool   { return true }
func (e *timeoutError) Temporary() bool { return true }

type AddrError struct {
	Err  string
	Addr string
}

func (e *AddrError) Error() string {
	if e == nil {
		return "<nil>"
	}
	s := e.Err
	if e.Addr != "" {
		s += " " + e.Addr
	}
	return s
}

func (e *AddrError) Temporary() bool {
	return false
}

func (e *AddrError) Timeout() bool {
	return false
}

type UnknownNetworkError string

func (e UnknownNetworkError) Error() string   { return "unknown network " + string(e) }
func (e UnknownNetworkError) Temporary() bool { return false }
func (e UnknownNetworkError) Timeout() bool   { return false }

type InvalidAddrError string

func (e InvalidAddrError) Error() string   { return string(e) }
func (e InvalidAddrError) Timeout() bool   { return false }
func (e InvalidAddrError) Temporary() bool { return false }

// DNSConfigError represents an error reading the machine's DNS configuration.
type DNSConfigError struct {
	Err error
}

func (e *DNSConfigError) Error() string {
	return "error reading DNS config: " + e.Err.Error()
}

func (e *DNSConfigError) Timeout() bool   { return false }
func (e *DNSConfigError) Temporary() bool { return false }

type writerOnly struct {
	io.Writer
}

// Fallback implementation of io.ReaderFrom's ReadFrom, when sendfile isn't
// applicable.
func genericReadFrom(w io.Writer, r io.Reader) (n int64, err error) {
	// Use wrapper to hide existing r.ReadFrom from io.Copy.
	return io.Copy(writerOnly{w}, r)
}

// Limit the number of concurrent cgo-using goroutines, because
// each will block an entire operating system thread. The usual culprit
// is resolving many DNS names in separate goroutines but the DNS
// server is not responding. Then the many lookups each use a different
// thread, and the system or the program runs out of threads.

var threadLimit = make(chan struct{}, 500)

// Using send for acquire is fine here because we are not using this
// to protect any memory. All we care about is the number of goroutines
// making calls at a time.

func acquireThread() {
	threadLimit <- struct{}{}
}

func releaseThread() {
	<-threadLimit
}
<|endoftext|>


<commit_before>package mem_constrained_buffer

import (
	"bytes"
	"io"
	"io/ioutil"
	"mime/multipart"
	"os"
)

const (
	DefaultMemorySize int64 = 1 << 18
	FilenamePrefix          = "mem-buf-"
)

type MemoryConstrainedBuffer struct {
	b             *bytes.Buffer
	tmpfile       string
	max           int64
	size          int64
	removeOnClose bool
	file          multipart.File
}

func NewMemoryConstrainedBuffer(maxMemory int64, removeOnClose bool) *MemoryConstrainedBuffer {
	return &MemoryConstrainedBuffer{
		max:           maxMemory,
		removeOnClose: removeOnClose,
	}
}

func NewDefaultMemoryConstrainedBuffer() *MemoryConstrainedBuffer {
	return NewMemoryConstrainedBuffer(DefaultMemorySize, true)
}

func (m *MemoryConstrainedBuffer) Write(p []byte) (int, error) {
	n, err := m.ReadFrom(bytes.NewReader(p))
	return int(n), err
}

func (m *MemoryConstrainedBuffer) ReadFrom(r io.Reader) (int64, error) {
	var (
		n   int64
		err error
		buf bytes.Buffer
	)

	for {
		n, err = io.CopyN(&buf, r, m.max+1)
		if err != nil && err != io.EOF {
			return 0, err
		}

		m.size += n

		if err == io.EOF {
			err = nil
			break
		}

		if n > m.max {
			// too big, write to disk and flush buffer
			file, err := ioutil.TempFile("", FilenamePrefix)
			if err != nil {
				return 0, err
			}
			defer file.Close()
			n, err = io.Copy(file, io.MultiReader(&buf, r))
			if err != nil {
				os.Remove(file.Name())
				return 0, err
			}
			m.tmpfile = file.Name()
			m.size = n
			break

		} else {
			m.max -= n
		}
	}

	if m.tmpfile == "" {
		m.b = &buf
	}

	return n, err
}

func (m *MemoryConstrainedBuffer) Len() int64 {
	return m.size
}

func (m *MemoryConstrainedBuffer) Remove() (err error) {
	if m.file == nil && m.tmpfile == "" {
		return nil
	}
	if m.file != nil {
		err = m.file.Close()
	}
	if m.tmpfile != "" {
		err = os.Remove(m.tmpfile)
	}
	return err
}

func (m *MemoryConstrainedBuffer) open() error {
	if m.file != nil {
		return nil
	}
	if b := m.b; b != nil {
		m.file = &sectionReadCloser{io.NewSectionReader(bytes.NewReader(m.b.Bytes()), 0, int64(b.Len()))}
		return nil
	}
	f, err := os.Open(m.tmpfile)
	m.file = f
	return err
}

func (m *MemoryConstrainedBuffer) Read(p []byte) (int, error) {
	if m.file == nil {
		if err := m.open(); err != nil {
			return 0, err
		}
	}
	return m.file.Read(p)
}

func (m *MemoryConstrainedBuffer) ReadAt(p []byte, off int64) (int, error) {
	if m.file == nil {
		if err := m.open(); err != nil {
			return 0, err
		}
	}
	return m.file.ReadAt(p, off)
}

func (m *MemoryConstrainedBuffer) Seek(offset int64, whence int) (int64, error) {
	if m.file == nil {
		if err := m.open(); err != nil {
			return 0, err
		}
	}
	return m.file.Seek(offset, whence)
}

func (m *MemoryConstrainedBuffer) Close() error {
	if m.removeOnClose {
		defer m.Remove()
	}
	if m.file == nil {
		return nil
	}
	return m.file.Close()
}

type sectionReadCloser struct {
	*io.SectionReader
}

func (rc sectionReadCloser) Close() error {
	return nil
}
<commit_msg>cleanup and refactoring<commit_after>package mem_constrained_buffer

import (
	"bytes"
	"io"
	"io/ioutil"
	"mime/multipart"
	"os"
)

const (
	DefaultMemorySize int64 = 1 << 18
	FilenamePrefix          = "mem-buf-"
)

type MemoryConstrainedBuffer struct {
	b             *bytes.Buffer
	tmpfile       string
	max           int64
	size          int64
	removeOnClose bool
	file          multipart.File
}

func NewMemoryConstrainedBuffer(maxMemory int64, removeOnClose bool) *MemoryConstrainedBuffer {
	return &MemoryConstrainedBuffer{
		max:           maxMemory,
		removeOnClose: removeOnClose,
		b:             new(bytes.Buffer),
	}
}

func NewDefaultMemoryConstrainedBuffer() *MemoryConstrainedBuffer {
	return NewMemoryConstrainedBuffer(DefaultMemorySize, true)
}

func (m *MemoryConstrainedBuffer) Write(p []byte) (int, error) {
	n, err := m.ReadFrom(bytes.NewReader(p))
	return int(n), err
}

func (m *MemoryConstrainedBuffer) open() error {
	if m.file != nil {
		return nil
	}
	if b := m.b; b != nil {
		m.file = &sectionReadCloser{io.NewSectionReader(bytes.NewReader(m.b.Bytes()), 0, int64(b.Len()))}
		return nil
	}
	f, err := os.Open(m.tmpfile)
	m.file = f
	return err
}

func (m *MemoryConstrainedBuffer) Read(p []byte) (int, error) {
	if err := m.open(); err != nil {
		return 0, err
	}
	n, err := m.file.Read(p)
	return n, err
}

func (m *MemoryConstrainedBuffer) ReadAt(p []byte, off int64) (int, error) {
	if m.file == nil {
		if err := m.open(); err != nil {
			return 0, err
		}
	}
	return m.file.ReadAt(p, off)
}

func (m *MemoryConstrainedBuffer) ReadFrom(r io.Reader) (int64, error) {
	var (
		n   int64
		err error
	)

	for {
		n, err = io.CopyN(m.b, r, m.max+1)
		if err != nil && err != io.EOF {
			return 0, err
		}

		m.size += n

		if err == io.EOF {
			err = nil
			break
		}

		if n > m.max {
			// too big, write to disk and flush buffer
			file, err := ioutil.TempFile("", FilenamePrefix)
			if err != nil {
				return 0, err
			}
			n, err = io.Copy(file, io.MultiReader(m.b, r))
			if err != nil {
				os.Remove(file.Name())
				return 0, err
			}
			m.b.Reset()
			m.tmpfile = file.Name()
			m.file = file
			m.file.Seek(0, 0)
			m.size = n
			break

		} else {
			m.max -= n
		}
	}

	return n, err
}

func (m *MemoryConstrainedBuffer) Len() int64 {
	return m.size
}

func (m *MemoryConstrainedBuffer) Seek(offset int64, whence int) (int64, error) {
	if err := m.open(); err != nil {
		return 0, err
	}
	return m.file.Seek(offset, whence)
}

func (m *MemoryConstrainedBuffer) Remove() (err error) {
	if m.file == nil && m.tmpfile == "" {
		return nil
	}
	if m.file != nil {
		err = m.file.Close()
	}
	if m.tmpfile != "" {
		err = os.Remove(m.tmpfile)
	}
	return err
}

func (m *MemoryConstrainedBuffer) Close() error {
	m.b.Reset()
	if m.file == nil {
		return nil
	}
	err := m.file.Close()
	if m.removeOnClose {
		err = m.Remove()
	}
	return err
}

type sectionReadCloser struct {
	*io.SectionReader
}

func (rc sectionReadCloser) Close() error {
	return nil
}
<|endoftext|>


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

package atomic

import (
	"unsafe"
)

// A Value provides an atomic load and store of a consistently typed value.
// The zero value for a Value returns nil from Load.
// Once Store has been called, a Value must not be copied.
//
// A Value must not be copied after first use.
type Value struct {
	v any
}

// ifaceWords is interface{} internal representation.
type ifaceWords struct {
	typ  unsafe.Pointer
	data unsafe.Pointer
}

// Load returns the value set by the most recent Store.
// It returns nil if there has been no call to Store for this Value.
func (v *Value) Load() (val any) {
	vp := (*ifaceWords)(unsafe.Pointer(v))
	typ := LoadPointer(&vp.typ)
	if typ == nil || typ == unsafe.Pointer(&firstStoreInProgress) {
		// First store not yet completed.
		return nil
	}
	data := LoadPointer(&vp.data)
	vlp := (*ifaceWords)(unsafe.Pointer(&val))
	vlp.typ = typ
	vlp.data = data
	return
}

var firstStoreInProgress byte

// Store sets the value of the Value to x.
// All calls to Store for a given Value must use values of the same concrete type.
// Store of an inconsistent type panics, as does Store(nil).
func (v *Value) Store(val any) {
	if val == nil {
		panic("sync/atomic: store of nil value into Value")
	}
	vp := (*ifaceWords)(unsafe.Pointer(v))
	vlp := (*ifaceWords)(unsafe.Pointer(&val))
	for {
		typ := LoadPointer(&vp.typ)
		if typ == nil {
			// Attempt to start first store.
			// Disable preemption so that other goroutines can use
			// active spin wait to wait for completion.
			runtime_procPin()
			if !CompareAndSwapPointer(&vp.typ, nil, unsafe.Pointer(&firstStoreInProgress)) {
				runtime_procUnpin()
				continue
			}
			// Complete first store.
			StorePointer(&vp.data, vlp.data)
			StorePointer(&vp.typ, vlp.typ)
			runtime_procUnpin()
			return
		}
		if typ == unsafe.Pointer(&firstStoreInProgress) {
			// First store in progress. Wait.
			// Since we disable preemption around the first store,
			// we can wait with active spinning.
			continue
		}
		// First store completed. Check type and overwrite data.
		if typ != vlp.typ {
			panic("sync/atomic: store of inconsistently typed value into Value")
		}
		StorePointer(&vp.data, vlp.data)
		return
	}
}

// Swap stores new into Value and returns the previous value. It returns nil if
// the Value is empty.
//
// All calls to Swap for a given Value must use values of the same concrete
// type. Swap of an inconsistent type panics, as does Swap(nil).
func (v *Value) Swap(new any) (old any) {
	if new == nil {
		panic("sync/atomic: swap of nil value into Value")
	}
	vp := (*ifaceWords)(unsafe.Pointer(v))
	np := (*ifaceWords)(unsafe.Pointer(&new))
	for {
		typ := LoadPointer(&vp.typ)
		if typ == nil {
			// Attempt to start first store.
			// Disable preemption so that other goroutines can use
			// active spin wait to wait for completion; and so that
			// GC does not see the fake type accidentally.
			runtime_procPin()
			if !CompareAndSwapPointer(&vp.typ, nil, unsafe.Pointer(&firstStoreInProgress)) {
				runtime_procUnpin()
				continue
			}
			// Complete first store.
			StorePointer(&vp.data, np.data)
			StorePointer(&vp.typ, np.typ)
			runtime_procUnpin()
			return nil
		}
		if typ == unsafe.Pointer(&firstStoreInProgress) {
			// First store in progress. Wait.
			// Since we disable preemption around the first store,
			// we can wait with active spinning.
			continue
		}
		// First store completed. Check type and overwrite data.
		if typ != np.typ {
			panic("sync/atomic: swap of inconsistently typed value into Value")
		}
		op := (*ifaceWords)(unsafe.Pointer(&old))
		op.typ, op.data = np.typ, SwapPointer(&vp.data, np.data)
		return old
	}
}

// CompareAndSwap executes the compare-and-swap operation for the Value.
//
// All calls to CompareAndSwap for a given Value must use values of the same
// concrete type. CompareAndSwap of an inconsistent type panics, as does
// CompareAndSwap(old, nil).
func (v *Value) CompareAndSwap(old, new any) (swapped bool) {
	if new == nil {
		panic("sync/atomic: compare and swap of nil value into Value")
	}
	vp := (*ifaceWords)(unsafe.Pointer(v))
	np := (*ifaceWords)(unsafe.Pointer(&new))
	op := (*ifaceWords)(unsafe.Pointer(&old))
	if op.typ != nil && np.typ != op.typ {
		panic("sync/atomic: compare and swap of inconsistently typed values")
	}
	for {
		typ := LoadPointer(&vp.typ)
		if typ == nil {
			if old != nil {
				return false
			}
			// Attempt to start first store.
			// Disable preemption so that other goroutines can use
			// active spin wait to wait for completion; and so that
			// GC does not see the fake type accidentally.
			runtime_procPin()
			if !CompareAndSwapPointer(&vp.typ, nil, unsafe.Pointer(&firstStoreInProgress)) {
				runtime_procUnpin()
				continue
			}
			// Complete first store.
			StorePointer(&vp.data, np.data)
			StorePointer(&vp.typ, np.typ)
			runtime_procUnpin()
			return true
		}
		if typ == unsafe.Pointer(&firstStoreInProgress) {
			// First store in progress. Wait.
			// Since we disable preemption around the first store,
			// we can wait with active spinning.
			continue
		}
		// First store completed. Check type and overwrite data.
		if typ != np.typ {
			panic("sync/atomic: compare and swap of inconsistently typed value into Value")
		}
		// Compare old and current via runtime equality check.
		// This allows value types to be compared, something
		// not offered by the package functions.
		// CompareAndSwapPointer below only ensures vp.data
		// has not changed since LoadPointer.
		data := LoadPointer(&vp.data)
		var i any
		(*ifaceWords)(unsafe.Pointer(&i)).typ = typ
		(*ifaceWords)(unsafe.Pointer(&i)).data = data
		if i != old {
			return false
		}
		return CompareAndSwapPointer(&vp.data, data, np.data)
	}
}

// Disable/enable preemption, implemented in runtime.
func runtime_procPin()
func runtime_procUnpin()
<commit_msg>sync/atomic: rename "ifaceWords" to "efaceWords"<commit_after>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package atomic

import (
	"unsafe"
)

// A Value provides an atomic load and store of a consistently typed value.
// The zero value for a Value returns nil from Load.
// Once Store has been called, a Value must not be copied.
//
// A Value must not be copied after first use.
type Value struct {
	v any
}

// efaceWords is interface{} internal representation.
type efaceWords struct {
	typ  unsafe.Pointer
	data unsafe.Pointer
}

// Load returns the value set by the most recent Store.
// It returns nil if there has been no call to Store for this Value.
func (v *Value) Load() (val any) {
	vp := (*efaceWords)(unsafe.Pointer(v))
	typ := LoadPointer(&vp.typ)
	if typ == nil || typ == unsafe.Pointer(&firstStoreInProgress) {
		// First store not yet completed.
		return nil
	}
	data := LoadPointer(&vp.data)
	vlp := (*efaceWords)(unsafe.Pointer(&val))
	vlp.typ = typ
	vlp.data = data
	return
}

var firstStoreInProgress byte

// Store sets the value of the Value to x.
// All calls to Store for a given Value must use values of the same concrete type.
// Store of an inconsistent type panics, as does Store(nil).
func (v *Value) Store(val any) {
	if val == nil {
		panic("sync/atomic: store of nil value into Value")
	}
	vp := (*efaceWords)(unsafe.Pointer(v))
	vlp := (*efaceWords)(unsafe.Pointer(&val))
	for {
		typ := LoadPointer(&vp.typ)
		if typ == nil {
			// Attempt to start first store.
			// Disable preemption so that other goroutines can use
			// active spin wait to wait for completion.
			runtime_procPin()
			if !CompareAndSwapPointer(&vp.typ, nil, unsafe.Pointer(&firstStoreInProgress)) {
				runtime_procUnpin()
				continue
			}
			// Complete first store.
			StorePointer(&vp.data, vlp.data)
			StorePointer(&vp.typ, vlp.typ)
			runtime_procUnpin()
			return
		}
		if typ == unsafe.Pointer(&firstStoreInProgress) {
			// First store in progress. Wait.
			// Since we disable preemption around the first store,
			// we can wait with active spinning.
			continue
		}
		// First store completed. Check type and overwrite data.
		if typ != vlp.typ {
			panic("sync/atomic: store of inconsistently typed value into Value")
		}
		StorePointer(&vp.data, vlp.data)
		return
	}
}

// Swap stores new into Value and returns the previous value. It returns nil if
// the Value is empty.
//
// All calls to Swap for a given Value must use values of the same concrete
// type. Swap of an inconsistent type panics, as does Swap(nil).
func (v *Value) Swap(new any) (old any) {
	if new == nil {
		panic("sync/atomic: swap of nil value into Value")
	}
	vp := (*efaceWords)(unsafe.Pointer(v))
	np := (*efaceWords)(unsafe.Pointer(&new))
	for {
		typ := LoadPointer(&vp.typ)
		if typ == nil {
			// Attempt to start first store.
			// Disable preemption so that other goroutines can use
			// active spin wait to wait for completion; and so that
			// GC does not see the fake type accidentally.
			runtime_procPin()
			if !CompareAndSwapPointer(&vp.typ, nil, unsafe.Pointer(&firstStoreInProgress)) {
				runtime_procUnpin()
				continue
			}
			// Complete first store.
			StorePointer(&vp.data, np.data)
			StorePointer(&vp.typ, np.typ)
			runtime_procUnpin()
			return nil
		}
		if typ == unsafe.Pointer(&firstStoreInProgress) {
			// First store in progress. Wait.
			// Since we disable preemption around the first store,
			// we can wait with active spinning.
			continue
		}
		// First store completed. Check type and overwrite data.
		if typ != np.typ {
			panic("sync/atomic: swap of inconsistently typed value into Value")
		}
		op := (*efaceWords)(unsafe.Pointer(&old))
		op.typ, op.data = np.typ, SwapPointer(&vp.data, np.data)
		return old
	}
}

// CompareAndSwap executes the compare-and-swap operation for the Value.
//
// All calls to CompareAndSwap for a given Value must use values of the same
// concrete type. CompareAndSwap of an inconsistent type panics, as does
// CompareAndSwap(old, nil).
func (v *Value) CompareAndSwap(old, new any) (swapped bool) {
	if new == nil {
		panic("sync/atomic: compare and swap of nil value into Value")
	}
	vp := (*efaceWords)(unsafe.Pointer(v))
	np := (*efaceWords)(unsafe.Pointer(&new))
	op := (*efaceWords)(unsafe.Pointer(&old))
	if op.typ != nil && np.typ != op.typ {
		panic("sync/atomic: compare and swap of inconsistently typed values")
	}
	for {
		typ := LoadPointer(&vp.typ)
		if typ == nil {
			if old != nil {
				return false
			}
			// Attempt to start first store.
			// Disable preemption so that other goroutines can use
			// active spin wait to wait for completion; and so that
			// GC does not see the fake type accidentally.
			runtime_procPin()
			if !CompareAndSwapPointer(&vp.typ, nil, unsafe.Pointer(&firstStoreInProgress)) {
				runtime_procUnpin()
				continue
			}
			// Complete first store.
			StorePointer(&vp.data, np.data)
			StorePointer(&vp.typ, np.typ)
			runtime_procUnpin()
			return true
		}
		if typ == unsafe.Pointer(&firstStoreInProgress) {
			// First store in progress. Wait.
			// Since we disable preemption around the first store,
			// we can wait with active spinning.
			continue
		}
		// First store completed. Check type and overwrite data.
		if typ != np.typ {
			panic("sync/atomic: compare and swap of inconsistently typed value into Value")
		}
		// Compare old and current via runtime equality check.
		// This allows value types to be compared, something
		// not offered by the package functions.
		// CompareAndSwapPointer below only ensures vp.data
		// has not changed since LoadPointer.
		data := LoadPointer(&vp.data)
		var i any
		(*efaceWords)(unsafe.Pointer(&i)).typ = typ
		(*efaceWords)(unsafe.Pointer(&i)).data = data
		if i != old {
			return false
		}
		return CompareAndSwapPointer(&vp.data, data, np.data)
	}
}

// Disable/enable preemption, implemented in runtime.
func runtime_procPin()
func runtime_procUnpin()
<|endoftext|>


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

// +build darwin dragonfly freebsd linux netbsd openbsd solaris

// Fork, exec, wait, etc.

package syscall

import (
	"runtime"
	"sync"
	"unsafe"
)

// Lock synchronizing creation of new file descriptors with fork.
//
// We want the child in a fork/exec sequence to inherit only the
// file descriptors we intend. To do that, we mark all file
// descriptors close-on-exec and then, in the child, explicitly
// unmark the ones we want the exec'ed program to keep.
// Unix doesn't make this easy: there is, in general, no way to
// allocate a new file descriptor close-on-exec. Instead you
// have to allocate the descriptor and then mark it close-on-exec.
// If a fork happens between those two events, the child's exec
// will inherit an unwanted file descriptor.
//
// This lock solves that race: the create new fd/mark close-on-exec
// operation is done holding ForkLock for reading, and the fork itself
// is done holding ForkLock for writing. At least, that's the idea.
// There are some complications.
//
// Some system calls that create new file descriptors can block
// for arbitrarily long times: open on a hung NFS server or named
// pipe, accept on a socket, and so on. We can't reasonably grab
// the lock across those operations.
//
// It is worse to inherit some file descriptors than others.
// If a non-malicious child accidentally inherits an open ordinary file,
// that's not a big deal. On the other hand, if a long-lived child
// accidentally inherits the write end of a pipe, then the reader
// of that pipe will not see EOF until that child exits, potentially
// causing the parent program to hang. This is a common problem
// in threaded C programs that use popen.
//
// Luckily, the file descriptors that are most important not to
// inherit are not the ones that can take an arbitrarily long time
// to create: pipe returns instantly, and the net package uses
// non-blocking I/O to accept on a listening socket.
// The rules for which file descriptor-creating operations use the
// ForkLock are as follows:
//
// 1) Pipe. Does not block. Use the ForkLock.
// 2) Socket. Does not block. Use the ForkLock.
// 3) Accept. If using non-blocking mode, use the ForkLock.
//             Otherwise, live with the race.
// 4) Open. Can block. Use O_CLOEXEC if available (Linux).
//             Otherwise, live with the race.
// 5) Dup. Does not block. Use the ForkLock.
//             On Linux, could use fcntl F_DUPFD_CLOEXEC
//             instead of the ForkLock, but only for dup(fd, -1).

var ForkLock sync.RWMutex

// StringSlicePtr converts a slice of strings to a slice of pointers
// to NUL-terminated byte arrays. If any string contains a NUL byte
// this function panics instead of returning an error.
//
// Deprecated: Use SlicePtrFromStrings instead.
func StringSlicePtr(ss []string) []*byte {
	bb := make([]*byte, len(ss)+1)
	for i := 0; i < len(ss); i++ {
		bb[i] = StringBytePtr(ss[i])
	}
	bb[len(ss)] = nil
	return bb
}

// SlicePtrFromStrings converts a slice of strings to a slice of
// pointers to NUL-terminated byte arrays. If any string contains
// a NUL byte, it returns (nil, EINVAL).
func SlicePtrFromStrings(ss []string) ([]*byte, error) {
	var err error
	bb := make([]*byte, len(ss)+1)
	for i := 0; i < len(ss); i++ {
		bb[i], err = BytePtrFromString(ss[i])
		if err != nil {
			return nil, err
		}
	}
	bb[len(ss)] = nil
	return bb, nil
}

func CloseOnExec(fd int) { fcntl(fd, F_SETFD, FD_CLOEXEC) }

func SetNonblock(fd int, nonblocking bool) (err error) {
	flag, err := fcntl(fd, F_GETFL, 0)
	if err != nil {
		return err
	}
	if nonblocking {
		flag |= O_NONBLOCK
	} else {
		flag &^= O_NONBLOCK
	}
	_, err = fcntl(fd, F_SETFL, flag)
	return err
}

// Credential holds user and group identities to be assumed
// by a child process started by StartProcess.
type Credential struct {
	Uid    uint32   // User ID.
	Gid    uint32   // Group ID.
	Groups []uint32 // Supplementary group IDs.
}

// ProcAttr holds attributes that will be applied to a new process started
// by StartProcess.
type ProcAttr struct {
	Dir   string    // Current working directory.
	Env   []string  // Environment.
	Files []uintptr // File descriptors.
	Sys   *SysProcAttr
}

var zeroProcAttr ProcAttr
var zeroSysProcAttr SysProcAttr

func forkExec(argv0 string, argv []string, attr *ProcAttr) (pid int, err error) {
	var p [2]int
	var n int
	var err1 Errno
	var wstatus WaitStatus

	if attr == nil {
		attr = &zeroProcAttr
	}
	sys := attr.Sys
	if sys == nil {
		sys = &zeroSysProcAttr
	}

	p[0] = -1
	p[1] = -1

	// Convert args to C form.
	argv0p, err := BytePtrFromString(argv0)
	if err != nil {
		return 0, err
	}
	argvp, err := SlicePtrFromStrings(argv)
	if err != nil {
		return 0, err
	}
	envvp, err := SlicePtrFromStrings(attr.Env)
	if err != nil {
		return 0, err
	}

	if (runtime.GOOS == "freebsd" || runtime.GOOS == "dragonfly") && len(argv[0]) > len(argv0) {
		argvp[0] = argv0p
	}

	var chroot *byte
	if sys.Chroot != "" {
		chroot, err = BytePtrFromString(sys.Chroot)
		if err != nil {
			return 0, err
		}
	}
	var dir *byte
	if attr.Dir != "" {
		dir, err = BytePtrFromString(attr.Dir)
		if err != nil {
			return 0, err
		}
	}

	// Acquire the fork lock so that no other threads
	// create new fds that are not yet close-on-exec
	// before we fork.
	ForkLock.Lock()

	// Allocate child status pipe close on exec.
	if err = forkExecPipe(p[:]); err != nil {
		goto error
	}

	// Kick off child.
	pid, err1 = forkAndExecInChild(argv0p, argvp, envvp, chroot, dir, attr, sys, p[1])
	if err1 != 0 {
		err = Errno(err1)
		goto error
	}
	ForkLock.Unlock()

	// Read child error status from pipe.
	Close(p[1])
	n, err = readlen(p[0], (*byte)(unsafe.Pointer(&err1)), int(unsafe.Sizeof(err1)))
	Close(p[0])
	if err != nil || n != 0 {
		if n == int(unsafe.Sizeof(err1)) {
			err = Errno(err1)
		}
		if err == nil {
			err = EPIPE
		}

		// Child failed; wait for it to exit, to make sure
		// the zombies don't accumulate.
		_, err1 := Wait4(pid, &wstatus, 0, nil)
		for err1 == EINTR {
			_, err1 = Wait4(pid, &wstatus, 0, nil)
		}
		return 0, err
	}

	// Read got EOF, so pipe closed on exec, so exec succeeded.
	return pid, nil

error:
	if p[0] >= 0 {
		Close(p[0])
		Close(p[1])
	}
	ForkLock.Unlock()
	return 0, err
}

// Combination of fork and exec, careful to be thread safe.
func ForkExec(argv0 string, argv []string, attr *ProcAttr) (pid int, err error) {
	return forkExec(argv0, argv, attr)
}

// StartProcess wraps ForkExec for package os.
func StartProcess(argv0 string, argv []string, attr *ProcAttr) (pid int, handle uintptr, err error) {
	pid, err = forkExec(argv0, argv, attr)
	return pid, 0, err
}

// Ordinary exec.
func Exec(argv0 string, argv []string, envv []string) (err error) {
	argv0p, err := BytePtrFromString(argv0)
	if err != nil {
		return err
	}
	argvp, err := SlicePtrFromStrings(argv)
	if err != nil {
		return err
	}
	envvp, err := SlicePtrFromStrings(envv)
	if err != nil {
		return err
	}
	_, _, err1 := RawSyscall(SYS_EXECVE,
		uintptr(unsafe.Pointer(argv0p)),
		uintptr(unsafe.Pointer(&argvp[0])),
		uintptr(unsafe.Pointer(&envvp[0])))
	return Errno(err1)
}
<commit_msg>syscall: document that Exec wraps execve(2)<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build darwin dragonfly freebsd linux netbsd openbsd solaris

// Fork, exec, wait, etc.

package syscall

import (
	"runtime"
	"sync"
	"unsafe"
)

// Lock synchronizing creation of new file descriptors with fork.
//
// We want the child in a fork/exec sequence to inherit only the
// file descriptors we intend. To do that, we mark all file
// descriptors close-on-exec and then, in the child, explicitly
// unmark the ones we want the exec'ed program to keep.
// Unix doesn't make this easy: there is, in general, no way to
// allocate a new file descriptor close-on-exec. Instead you
// have to allocate the descriptor and then mark it close-on-exec.
// If a fork happens between those two events, the child's exec
// will inherit an unwanted file descriptor.
//
// This lock solves that race: the create new fd/mark close-on-exec
// operation is done holding ForkLock for reading, and the fork itself
// is done holding ForkLock for writing. At least, that's the idea.
// There are some complications.
//
// Some system calls that create new file descriptors can block
// for arbitrarily long times: open on a hung NFS server or named
// pipe, accept on a socket, and so on. We can't reasonably grab
// the lock across those operations.
//
// It is worse to inherit some file descriptors than others.
// If a non-malicious child accidentally inherits an open ordinary file,
// that's not a big deal. On the other hand, if a long-lived child
// accidentally inherits the write end of a pipe, then the reader
// of that pipe will not see EOF until that child exits, potentially
// causing the parent program to hang. This is a common problem
// in threaded C programs that use popen.
//
// Luckily, the file descriptors that are most important not to
// inherit are not the ones that can take an arbitrarily long time
// to create: pipe returns instantly, and the net package uses
// non-blocking I/O to accept on a listening socket.
// The rules for which file descriptor-creating operations use the
// ForkLock are as follows:
//
// 1) Pipe. Does not block. Use the ForkLock.
// 2) Socket. Does not block. Use the ForkLock.
// 3) Accept. If using non-blocking mode, use the ForkLock.
//             Otherwise, live with the race.
// 4) Open. Can block. Use O_CLOEXEC if available (Linux).
//             Otherwise, live with the race.
// 5) Dup. Does not block. Use the ForkLock.
//             On Linux, could use fcntl F_DUPFD_CLOEXEC
//             instead of the ForkLock, but only for dup(fd, -1).

var ForkLock sync.RWMutex

// StringSlicePtr converts a slice of strings to a slice of pointers
// to NUL-terminated byte arrays. If any string contains a NUL byte
// this function panics instead of returning an error.
//
// Deprecated: Use SlicePtrFromStrings instead.
func StringSlicePtr(ss []string) []*byte {
	bb := make([]*byte, len(ss)+1)
	for i := 0; i < len(ss); i++ {
		bb[i] = StringBytePtr(ss[i])
	}
	bb[len(ss)] = nil
	return bb
}

// SlicePtrFromStrings converts a slice of strings to a slice of
// pointers to NUL-terminated byte arrays. If any string contains
// a NUL byte, it returns (nil, EINVAL).
func SlicePtrFromStrings(ss []string) ([]*byte, error) {
	var err error
	bb := make([]*byte, len(ss)+1)
	for i := 0; i < len(ss); i++ {
		bb[i], err = BytePtrFromString(ss[i])
		if err != nil {
			return nil, err
		}
	}
	bb[len(ss)] = nil
	return bb, nil
}

func CloseOnExec(fd int) { fcntl(fd, F_SETFD, FD_CLOEXEC) }

func SetNonblock(fd int, nonblocking bool) (err error) {
	flag, err := fcntl(fd, F_GETFL, 0)
	if err != nil {
		return err
	}
	if nonblocking {
		flag |= O_NONBLOCK
	} else {
		flag &^= O_NONBLOCK
	}
	_, err = fcntl(fd, F_SETFL, flag)
	return err
}

// Credential holds user and group identities to be assumed
// by a child process started by StartProcess.
type Credential struct {
	Uid    uint32   // User ID.
	Gid    uint32   // Group ID.
	Groups []uint32 // Supplementary group IDs.
}

// ProcAttr holds attributes that will be applied to a new process started
// by StartProcess.
type ProcAttr struct {
	Dir   string    // Current working directory.
	Env   []string  // Environment.
	Files []uintptr // File descriptors.
	Sys   *SysProcAttr
}

var zeroProcAttr ProcAttr
var zeroSysProcAttr SysProcAttr

func forkExec(argv0 string, argv []string, attr *ProcAttr) (pid int, err error) {
	var p [2]int
	var n int
	var err1 Errno
	var wstatus WaitStatus

	if attr == nil {
		attr = &zeroProcAttr
	}
	sys := attr.Sys
	if sys == nil {
		sys = &zeroSysProcAttr
	}

	p[0] = -1
	p[1] = -1

	// Convert args to C form.
	argv0p, err := BytePtrFromString(argv0)
	if err != nil {
		return 0, err
	}
	argvp, err := SlicePtrFromStrings(argv)
	if err != nil {
		return 0, err
	}
	envvp, err := SlicePtrFromStrings(attr.Env)
	if err != nil {
		return 0, err
	}

	if (runtime.GOOS == "freebsd" || runtime.GOOS == "dragonfly") && len(argv[0]) > len(argv0) {
		argvp[0] = argv0p
	}

	var chroot *byte
	if sys.Chroot != "" {
		chroot, err = BytePtrFromString(sys.Chroot)
		if err != nil {
			return 0, err
		}
	}
	var dir *byte
	if attr.Dir != "" {
		dir, err = BytePtrFromString(attr.Dir)
		if err != nil {
			return 0, err
		}
	}

	// Acquire the fork lock so that no other threads
	// create new fds that are not yet close-on-exec
	// before we fork.
	ForkLock.Lock()

	// Allocate child status pipe close on exec.
	if err = forkExecPipe(p[:]); err != nil {
		goto error
	}

	// Kick off child.
	pid, err1 = forkAndExecInChild(argv0p, argvp, envvp, chroot, dir, attr, sys, p[1])
	if err1 != 0 {
		err = Errno(err1)
		goto error
	}
	ForkLock.Unlock()

	// Read child error status from pipe.
	Close(p[1])
	n, err = readlen(p[0], (*byte)(unsafe.Pointer(&err1)), int(unsafe.Sizeof(err1)))
	Close(p[0])
	if err != nil || n != 0 {
		if n == int(unsafe.Sizeof(err1)) {
			err = Errno(err1)
		}
		if err == nil {
			err = EPIPE
		}

		// Child failed; wait for it to exit, to make sure
		// the zombies don't accumulate.
		_, err1 := Wait4(pid, &wstatus, 0, nil)
		for err1 == EINTR {
			_, err1 = Wait4(pid, &wstatus, 0, nil)
		}
		return 0, err
	}

	// Read got EOF, so pipe closed on exec, so exec succeeded.
	return pid, nil

error:
	if p[0] >= 0 {
		Close(p[0])
		Close(p[1])
	}
	ForkLock.Unlock()
	return 0, err
}

// Combination of fork and exec, careful to be thread safe.
func ForkExec(argv0 string, argv []string, attr *ProcAttr) (pid int, err error) {
	return forkExec(argv0, argv, attr)
}

// StartProcess wraps ForkExec for package os.
func StartProcess(argv0 string, argv []string, attr *ProcAttr) (pid int, handle uintptr, err error) {
	pid, err = forkExec(argv0, argv, attr)
	return pid, 0, err
}

// Exec invokes the execve(2) system call.
func Exec(argv0 string, argv []string, envv []string) (err error) {
	argv0p, err := BytePtrFromString(argv0)
	if err != nil {
		return err
	}
	argvp, err := SlicePtrFromStrings(argv)
	if err != nil {
		return err
	}
	envvp, err := SlicePtrFromStrings(envv)
	if err != nil {
		return err
	}
	_, _, err1 := RawSyscall(SYS_EXECVE,
		uintptr(unsafe.Pointer(argv0p)),
		uintptr(unsafe.Pointer(&argvp[0])),
		uintptr(unsafe.Pointer(&envvp[0])))
	return Errno(err1)
}
<|endoftext|>


<commit_before>package ssh

import (
	"bytes"
	"fmt"
	"github.com/lunixbochs/vtclean"
	"github.com/pkg/errors"
	"github.com/viant/toolbox"
	"golang.org/x/crypto/ssh"
	"io"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

//ErrTerminated - command session terminated
var ErrTerminated = errors.New("terminate")

const defaultShell = "/bin/bash"

const (
	drainTimeoutMs       = 10
	defaultTimeoutMs     = 5000
	initTimeoutMs        = 200
	defaultTickFrequency = 100
)

//Listener represent command listener (it will send stdout fragments as thier being available on stdout)
type Listener func(stdout string, hasMore bool)

//MultiCommandSession represents a multi command session
type MultiCommandSession interface {
	Run(command string, listener Listener, timeoutMs int, terminators ...string) (string, error)

	ShellPrompt() string

	System() string

	Reconnect() error

	Close()
}

//multiCommandSession represents a multi command session
//a new command are send vi stdin
type multiCommandSession struct {
	service            *service
	config             *SessionConfig
	replayCommands     *ReplayCommands
	recordSession      bool
	session            *ssh.Session
	stdOutput          chan string
	stdError           chan string
	stdInput           io.WriteCloser
	promptSequence     string
	shellPrompt        string
	escapedShellPrompt string
	system             string
	running            int32
}

func (s *multiCommandSession) Run(command string, listener Listener, timeoutMs int, terminators ...string) (string, error) {
	if atomic.LoadInt32(&s.running) == 0 {
		return "", ErrTerminated
	}
	s.drainStdout()
	var stdin = command + "\n"
	_, err := s.stdInput.Write([]byte(stdin))
	if err != nil {
		return "", fmt.Errorf("failed to execute command: %v, err: %v", command, err)
	}
	var output string
	output, _, err = s.readResponse(timeoutMs, listener, terminators...)
	if s.recordSession {
		s.replayCommands.Register(stdin, output)
	}
	return output, err
}

//ShellPrompt returns a shell prompt
func (s *multiCommandSession) ShellPrompt() string {
	return s.shellPrompt
}

//System returns a system name
func (s *multiCommandSession) System() string {
	return s.system
}

//Close closes the session with its resources
func (s *multiCommandSession) Close() {
	atomic.StoreInt32(&s.running, 0)
	s.stdInput.Close()
	if s.session != nil {
		s.session.Close()
	}

}

func (s *multiCommandSession) closeIfError(err error) bool {
	if err != nil {
		s.Close()
		return true
	}
	return false
}

func (s *multiCommandSession) start(shell string) (output string, err error) {
	var reader, errReader io.Reader
	reader, err = s.session.StdoutPipe()
	if err != nil {
		return "", err
	}
	errReader, err = s.session.StderrPipe()
	if err != nil {
		return "", err
	}

	waitGroup := sync.WaitGroup{}
	waitGroup.Add(2)
	go func() {
		waitGroup.Done()
		s.copy(reader, s.stdOutput)
	}()
	go func() {
		waitGroup.Done()
		s.copy(errReader, s.stdError)
	}()

	if shell == "" {
		shell = defaultShell
	}
	waitGroup.Wait()
	err = s.session.Start(shell)
	if err != nil {
		return "", err
	}
	output, _, err = s.readResponse(drainTimeoutMs, nil)
	return output, err
}

//copy copy data from reader to channel
func (s *multiCommandSession) copy(reader io.Reader, out chan string) {
	var written int64 = 0
	buf := make([]byte, 128*1024)
	var err error
	var bytesRead int
	for {
		writer := new(bytes.Buffer)
		if atomic.LoadInt32(&s.running) == 0 {
			return
		}
		bytesRead, err = reader.Read(buf)
		if bytesRead > 0 {
			bytesWritten, writeError := writer.Write(buf[:bytesRead])
			if s.closeIfError(writeError) {
				return
			}
			if bytesWritten > 0 {
				written += int64(bytesWritten)
			}

			if bytesRead != bytesWritten {
				if s.closeIfError(io.ErrShortWrite) {
					return
				}
			}
			out <- string(writer.Bytes())
		}

		if s.closeIfError(err) {
			return
		}
	}
}

func escapeInput(input string) string {
	input = vtclean.Clean(input, false)
	if input == "" {
		return input
	}
	return strings.Trim(input, "\n\r\t ")
}

func (s *multiCommandSession) Reconnect() (err error) {
	atomic.StoreInt32(&s.running, 1)
	s.service.Reconnect()
	s.session, err = s.service.client.NewSession()
	defer func() {
		if err != nil {
			s.service.client.Close()
		}
	}()

	if err != nil {
		return err
	}

	return s.init()
}

func (s *multiCommandSession) hasPrompt(input string) bool {
	escapedInput := escapeInput(input)
	var shellPrompt = s.shellPrompt
	if shellPrompt == "" {
		shellPrompt = "$"
	}
	if s.escapedShellPrompt == "" && s.shellPrompt != "" {
		s.escapedShellPrompt = escapeInput(s.shellPrompt)
	}

	if s.escapedShellPrompt != "" && strings.HasSuffix(escapedInput, s.escapedShellPrompt) || strings.HasSuffix(input, s.shellPrompt) {
		return true
	}
	return false
}

func (s *multiCommandSession) hasTerminator(input string, terminators ...string) bool {
	escapedInput := escapeInput(input)
	input = escapedInput
	for _, candidate := range terminators {
		candidateLen := len(candidate)
		if candidateLen == 0 {
			continue
		}
		if candidate[0:1] == "^" && strings.HasPrefix(input, candidate[1:]) {
			return true
		}
		if candidate[candidateLen-1:] == "$" && strings.HasSuffix(input, candidate[:candidateLen-1]) {
			return true
		}
		if strings.Contains(input, candidate) {
			return true
		}
	}
	return false
}

func (s *multiCommandSession) removePromptIfNeeded(stdout string) string {
	if strings.Contains(stdout, s.shellPrompt) {
		stdout = strings.Replace(stdout, s.shellPrompt, "", 1)
		stdout = strings.Replace(stdout, "\r", "", len(stdout))
		var lines = []string{}
		for _, line := range strings.Split(stdout, "\n") {
			if strings.TrimSpace(line) == "" {
				continue
			}
			lines = append(lines, line)
		}
		stdout = strings.Join(lines, "\r\n")
	}
	return stdout
}

func (s *multiCommandSession) readResponse(timeoutMs int, listener Listener, terminators ...string) (out string, has bool, err error) {
	if timeoutMs == 0 {
		timeoutMs = defaultTimeoutMs
	}
	defer func() {
		if listener != nil {
			listener("", false)
		}
	}()

	var done int32
	defer atomic.StoreInt32(&done, 1)
	var errOut string
	var hasOutput bool

	var waitTimeMs = 0
	var tickFrequencyMs = defaultTickFrequency
	if tickFrequencyMs > timeoutMs {
		tickFrequencyMs = timeoutMs
	}
	var timeoutDuration = time.Duration(tickFrequencyMs) * time.Millisecond

	var hasPrompt, hasTerminator bool

outer:
	for {
		select {
		case o := <-s.stdOutput:
			if len(o) > 0 {
				waitTimeMs = 0
				if listener != nil {
					listener(s.removePromptIfNeeded(o), true)
				}
			}
			out += o
			hasPrompt = s.hasPrompt(out)
			hasTerminator = s.hasTerminator(out, terminators...)
			if (hasPrompt || hasTerminator) && len(s.stdOutput) == 0 {
				break outer
			}
		case e := <-s.stdError:
			errOut += e
			if listener != nil {
				listener(s.removePromptIfNeeded(e), true)
			}
			hasPrompt = s.hasPrompt(errOut)
			hasTerminator = s.hasTerminator(errOut, terminators...)
			if (hasPrompt || hasTerminator) && len(s.stdOutput) == 0 {
				break outer
			}
		case <-time.After(timeoutDuration):
			waitTimeMs += tickFrequencyMs
			if waitTimeMs >= timeoutMs {
				break outer
			}
		}
	}
	if hasTerminator {
		s.drainStdout()
	}
	if errOut != "" {
		err = errors.New(errOut)
	}

	if len(out) > 0 {
		hasOutput = true
		out = s.removePromptIfNeeded(out)
	}
	return out, hasOutput, err
}

func (s *multiCommandSession) drainStdout() {
	//read any outstanding output
	for {
		_, has, _ := s.readResponse(drainTimeoutMs, nil, "")
		if !has {
			return
		}
	}
}

func (s *multiCommandSession) shellInit() (err error) {
	if s.promptSequence != "" {
		if _, err = s.Run(s.promptSequence, nil, initTimeoutMs); err != nil {
			return err
		}
	}

	var ts = toolbox.AsString(time.Now().UnixNano())
	s.promptSequence = "PS1=\"" + ts + "\\$\""
	_, err = s.Run(s.promptSequence, nil, initTimeoutMs)
	if s.closeIfError(err) {
		return err
	}
	s.shellPrompt = ts + "$"
	s.escapedShellPrompt = escapeInput(s.shellPrompt)

	s.drainStdout()
	for i := 0; i < 3; i++ {
		s.system, err = s.Run("uname -s", nil, defaultTimeoutMs)
		s.system = strings.ToLower(strings.TrimSpace(s.system))
		if s.system != "" && !strings.Contains(s.system, "$") {
			break
		}

	}
	return nil
}

func (s *multiCommandSession) init() (err error) {
	s.session, err = s.service.client.NewSession()
	defer func() {
		if err != nil {
			s.service.client.Close()
		}
	}()
	s.stdOutput = make(chan string)
	s.stdError = make(chan string)
	for k, v := range s.config.EnvVariables {
		err = s.session.Setenv(k, v)
		if err != nil {
			return err
		}
	}
	modes := ssh.TerminalModes{
		ssh.ECHO:          0,     // disable echoing
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}

	if err := s.session.RequestPty(s.config.Term, s.config.Rows, s.config.Columns, modes); err != nil {
		return err
	}

	if s.stdInput, err = s.session.StdinPipe(); err != nil {
		return err
	}

	_, err = s.start(s.config.Shell)
	if s.closeIfError(err) {
		return err
	}
	return s.shellInit()
}

func newMultiCommandSession(service *service, config *SessionConfig, replayCommands *ReplayCommands, recordSession bool) (MultiCommandSession, error) {
	if config == nil {
		config = &SessionConfig{}
	}
	config.applyDefault()

	result := &multiCommandSession{
		service:        service,
		config:         config,
		running:        1,
		recordSession:  recordSession,
		replayCommands: replayCommands,
	}
	return result, result.init()
}
<commit_msg>added init wait for slow terminals<commit_after>package ssh

import (
	"bytes"
	"fmt"
	"github.com/lunixbochs/vtclean"
	"github.com/pkg/errors"
	"github.com/viant/toolbox"
	"golang.org/x/crypto/ssh"
	"io"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

//ErrTerminated - command session terminated
var ErrTerminated = errors.New("terminate")

const defaultShell = "/bin/bash"

const (
	drainTimeoutMs       = 10
	defaultTimeoutMs     = 10000
	initTimeoutMs        = 200
	defaultTickFrequency = 100
)

//Listener represent command listener (it will send stdout fragments as thier being available on stdout)
type Listener func(stdout string, hasMore bool)

//MultiCommandSession represents a multi command session
type MultiCommandSession interface {
	Run(command string, listener Listener, timeoutMs int, terminators ...string) (string, error)

	ShellPrompt() string

	System() string

	Reconnect() error

	Close()
}

//multiCommandSession represents a multi command session
//a new command are send vi stdin
type multiCommandSession struct {
	service            *service
	config             *SessionConfig
	replayCommands     *ReplayCommands
	recordSession      bool
	session            *ssh.Session
	stdOutput          chan string
	stdError           chan string
	stdInput           io.WriteCloser
	promptSequence     string
	shellPrompt        string
	escapedShellPrompt string
	system             string
	running            int32
}

func (s *multiCommandSession) Run(command string, listener Listener, timeoutMs int, terminators ...string) (string, error) {
	if atomic.LoadInt32(&s.running) == 0 {
		return "", ErrTerminated
	}
	s.drainStdout()
	var stdin = command + "\n"
	_, err := s.stdInput.Write([]byte(stdin))
	if err != nil {
		return "", fmt.Errorf("failed to execute command: %v, err: %v", command, err)
	}
	var output string
	output, _, err = s.readResponse(timeoutMs, listener, terminators...)
	if s.recordSession {
		s.replayCommands.Register(stdin, output)
	}
	return output, err
}

//ShellPrompt returns a shell prompt
func (s *multiCommandSession) ShellPrompt() string {
	return s.shellPrompt
}

//System returns a system name
func (s *multiCommandSession) System() string {
	return s.system
}

//Close closes the session with its resources
func (s *multiCommandSession) Close() {
	atomic.StoreInt32(&s.running, 0)
	s.stdInput.Close()
	if s.session != nil {
		s.session.Close()
	}

}

func (s *multiCommandSession) closeIfError(err error) bool {
	if err != nil {
		s.Close()
		return true
	}
	return false
}

func (s *multiCommandSession) start(shell string) (output string, err error) {
	var reader, errReader io.Reader
	reader, err = s.session.StdoutPipe()
	if err != nil {
		return "", err
	}
	errReader, err = s.session.StderrPipe()
	if err != nil {
		return "", err
	}

	waitGroup := sync.WaitGroup{}
	waitGroup.Add(2)
	go func() {
		waitGroup.Done()
		s.copy(reader, s.stdOutput)
	}()
	go func() {
		waitGroup.Done()
		s.copy(errReader, s.stdError)
	}()

	if shell == "" {
		shell = defaultShell
	}
	waitGroup.Wait()
	err = s.session.Start(shell)
	if err != nil {
		return "", err
	}
	output, _, err = s.readResponse(drainTimeoutMs, nil)
	return output, err
}

//copy copy data from reader to channel
func (s *multiCommandSession) copy(reader io.Reader, out chan string) {
	var written int64 = 0
	buf := make([]byte, 128*1024)
	var err error
	var bytesRead int
	for {
		writer := new(bytes.Buffer)
		if atomic.LoadInt32(&s.running) == 0 {
			return
		}
		bytesRead, err = reader.Read(buf)
		if bytesRead > 0 {
			bytesWritten, writeError := writer.Write(buf[:bytesRead])
			if s.closeIfError(writeError) {
				return
			}
			if bytesWritten > 0 {
				written += int64(bytesWritten)
			}

			if bytesRead != bytesWritten {
				if s.closeIfError(io.ErrShortWrite) {
					return
				}
			}
			out <- string(writer.Bytes())
		}

		if s.closeIfError(err) {
			return
		}
	}
}

func escapeInput(input string) string {
	input = vtclean.Clean(input, false)
	if input == "" {
		return input
	}
	return strings.Trim(input, "\n\r\t ")
}

func (s *multiCommandSession) Reconnect() (err error) {
	atomic.StoreInt32(&s.running, 1)
	s.service.Reconnect()
	s.session, err = s.service.client.NewSession()
	defer func() {
		if err != nil {
			s.service.client.Close()
		}
	}()

	if err != nil {
		return err
	}

	return s.init()
}

func (s *multiCommandSession) hasPrompt(input string) bool {
	escapedInput := escapeInput(input)
	var shellPrompt = s.shellPrompt
	if shellPrompt == "" {
		shellPrompt = "$"
	}
	if s.escapedShellPrompt == "" && s.shellPrompt != "" {
		s.escapedShellPrompt = escapeInput(s.shellPrompt)
	}

	if s.escapedShellPrompt != "" && strings.HasSuffix(escapedInput, s.escapedShellPrompt) || strings.HasSuffix(input, s.shellPrompt) {
		return true
	}
	return false
}

func (s *multiCommandSession) hasTerminator(input string, terminators ...string) bool {
	escapedInput := escapeInput(input)
	input = escapedInput
	for _, candidate := range terminators {
		candidateLen := len(candidate)
		if candidateLen == 0 {
			continue
		}
		if candidate[0:1] == "^" && strings.HasPrefix(input, candidate[1:]) {
			return true
		}
		if candidate[candidateLen-1:] == "$" && strings.HasSuffix(input, candidate[:candidateLen-1]) {
			return true
		}
		if strings.Contains(input, candidate) {
			return true
		}
	}
	return false
}

func (s *multiCommandSession) removePromptIfNeeded(stdout string) string {
	if strings.Contains(stdout, s.shellPrompt) {
		stdout = strings.Replace(stdout, s.shellPrompt, "", 1)
		stdout = strings.Replace(stdout, "\r", "", len(stdout))
		var lines = []string{}
		for _, line := range strings.Split(stdout, "\n") {
			if strings.TrimSpace(line) == "" {
				continue
			}
			lines = append(lines, line)
		}
		stdout = strings.Join(lines, "\r\n")
	}
	return stdout
}

func (s *multiCommandSession) readResponse(timeoutMs int, listener Listener, terminators ...string) (out string, has bool, err error) {
	if timeoutMs == 0 {
		timeoutMs = defaultTimeoutMs
	}
	defer func() {
		if listener != nil {
			listener("", false)
		}
	}()

	var done int32
	defer atomic.StoreInt32(&done, 1)
	var errOut string
	var hasOutput bool

	var waitTimeMs = 0
	var tickFrequencyMs = defaultTickFrequency
	if tickFrequencyMs > timeoutMs {
		tickFrequencyMs = timeoutMs
	}
	var timeoutDuration = time.Duration(tickFrequencyMs) * time.Millisecond

	var hasPrompt, hasTerminator bool

outer:
	for {
		select {
		case o := <-s.stdOutput:
			if len(o) > 0 {
				waitTimeMs = 0
				if listener != nil {
					listener(s.removePromptIfNeeded(o), true)
				}
			}
			out += o
			hasPrompt = s.hasPrompt(out)
			hasTerminator = s.hasTerminator(out, terminators...)
			if (hasPrompt || hasTerminator) && len(s.stdOutput) == 0 {
				break outer
			}
		case e := <-s.stdError:
			errOut += e
			if listener != nil {
				listener(s.removePromptIfNeeded(e), true)
			}
			hasPrompt = s.hasPrompt(errOut)
			hasTerminator = s.hasTerminator(errOut, terminators...)
			if (hasPrompt || hasTerminator) && len(s.stdOutput) == 0 {
				break outer
			}
		case <-time.After(timeoutDuration):
			waitTimeMs += tickFrequencyMs
			if waitTimeMs >= timeoutMs {
				break outer
			}
		}
	}
	if hasTerminator {
		s.drainStdout()
	}
	if errOut != "" {
		err = errors.New(errOut)
	}

	if len(out) > 0 {
		hasOutput = true
		out = s.removePromptIfNeeded(out)
	}
	return out, hasOutput, err
}

func (s *multiCommandSession) drainStdout() {
	//read any outstanding output
	for {
		_, has, _ := s.readResponse(drainTimeoutMs, nil, "")
		if !has {
			return
		}
	}
}

func waitTimeout(wg *sync.WaitGroup, timeout time.Duration) bool {
	c := make(chan bool)
	go func() {
		defer close(c)
		wg.Wait()
		c <- true
	}()
	select {
	case <-c:
		return false // completed normally
	case <-time.After(timeout):
		return true // timed out
	}
}

func (s *multiCommandSession) shellInit() (err error) {
	if s.promptSequence != "" {
		if _, err = s.Run(s.promptSequence, nil, initTimeoutMs); err != nil {
			return err
		}
	}

	var ts = toolbox.AsString(time.Now().UnixNano())
	var waitGroup = &sync.WaitGroup{}
	waitGroup.Add(1)

	s.promptSequence = "PS1=\"" + ts + "\\$\""
	s.shellPrompt = ts + "$"
	s.escapedShellPrompt = escapeInput(s.shellPrompt)
	var listener Listener
	listener = func(stdout string, hasMore bool) {
		if ! hasMore {
			waitGroup.Done()
		}
	}
	_, err = s.Run("", listener, initTimeoutMs)
	waitTimeout(waitGroup, 60*time.Second)
	s.drainStdout()
	_, err = s.Run(s.promptSequence,  nil, defaultTimeoutMs)
	if s.closeIfError(err) {
		return err
	}
	for i := 0; i < 3; i++ {
		s.system, err = s.Run("uname -s", nil, initTimeoutMs)
		s.system = strings.ToLower(strings.TrimSpace(s.system))
		if strings.TrimSpace(s.system) != "" && !strings.Contains(s.system, "$") {
			break
		}
	}
	s.drainStdout()
	return nil
}

func (s *multiCommandSession) init() (err error) {
	s.session, err = s.service.client.NewSession()
	defer func() {
		if err != nil {
			s.service.client.Close()
		}
	}()
	s.stdOutput = make(chan string)
	s.stdError = make(chan string)
	for k, v := range s.config.EnvVariables {
		err = s.session.Setenv(k, v)
		if err != nil {
			return err
		}
	}
	modes := ssh.TerminalModes{
		ssh.ECHO:          0,     // disable echoing
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}

	if err := s.session.RequestPty(s.config.Term, s.config.Rows, s.config.Columns, modes); err != nil {
		return err
	}

	if s.stdInput, err = s.session.StdinPipe(); err != nil {
		return err
	}

	_, err = s.start(s.config.Shell)
	if s.closeIfError(err) {
		return err
	}
	return s.shellInit()
}

func newMultiCommandSession(service *service, config *SessionConfig, replayCommands *ReplayCommands, recordSession bool) (MultiCommandSession, error) {
	if config == nil {
		config = &SessionConfig{}
	}
	config.applyDefault()

	result := &multiCommandSession{
		service:        service,
		config:         config,
		running:        1,
		recordSession:  recordSession,
		replayCommands: replayCommands,
	}
	return result, result.init()
}
<|endoftext|>