<commit_before>package ssh

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

const defaultShell = "/bin/bash"

const defautTimeoutMs = 5000

//MultiCommandSession represents a multi command session
//a new command are send vi stdin
type MultiCommandSession struct {
	session     *ssh.Session
	stdOutput   chan string
	stdError    chan string
	stdInput    io.WriteCloser
	ShellPrompt string
	running     int32
}

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

func (s *MultiCommandSession) init(shell string) (string, error) {
	reader, err := s.session.StdoutPipe()
	if err != nil {
		return "", err
	}
	go s.drain(reader, s.stdOutput)

	errReader, err := s.session.StderrPipe()
	if err != nil {
		return "", err
	}
	go s.drain(errReader, s.stdError)
	if shell == "" {
		shell = defaultShell
	}
	err = s.session.Start(shell)
	if err != nil {
		return "", err
	}
	return s.readResponse(defautTimeoutMs)
}

func (s *MultiCommandSession) drain(reader io.Reader, out chan string) {
	var written int64 = 0
	buf := make([]byte, 128*1024)
	for {
		writter := new(bytes.Buffer)
		if atomic.LoadInt32(&s.running) == 0 {
			return
		}

		bytesRead, readError := reader.Read(buf)
		if bytesRead > 0 {
			bytesWritten, writeError := writter.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(writter.Bytes())
		}
		if s.closeIfError(readError) {
			return
		}

	}
}

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

//Close closes the session with its resources
func (s *MultiCommandSession) Close() {
	atomic.StoreInt32(&s.running, 0)
	s.stdInput.Close()
	s.session.Close()

}

func (s *MultiCommandSession) readResponse(timeoutMs int, terminators ...string) (out string, err error) {
	if timeoutMs == 0 {
		timeoutMs = defautTimeoutMs
	}
	if len(terminators) == 0 {
		if s.ShellPrompt == "" {
			terminators = []string{s.ShellPrompt + "$"}
		} else {
			terminators = []string{"$ $"}
		}
	}
	var done int32
	defer atomic.StoreInt32(&done, 1)
	var errOut string
outer:
	for {
		select {

		case o := <-s.stdOutput:
			out += o
			if hasTerminator(out, terminators...) && len(s.stdOutput) == 0 {
				break outer
			}
		case e := <-s.stdError:
			errOut += e
			if hasTerminator(errOut, terminators...) && len(s.stdOutput) == 0 {
				break outer
			}

		case <-time.After(time.Duration(timeoutMs) * time.Millisecond):
			break outer
		}
	}
	if errOut != "" {
		err = errors.New(errOut)
	}
	if len(out) > 0 {
		index := strings.LastIndex(out, "\r\n"+s.ShellPrompt)
		if index > 0 {
			out = string(out[:index])
		}
	}
	return out, err
}

func (s *MultiCommandSession) Run(command string, timeoutMs int, terminators ...string) (string, error) {
	_, err := s.stdInput.Write([]byte(command + "\n"))
	if err != nil {
		return "", fmt.Errorf("Failed to execute command: %v, err: %v", command, err)
	}
	return s.readResponse(timeoutMs, terminators...)
}

func newMultiCommandSession(client *ssh.Client, config *SessionConfig) (result *MultiCommandSession, err error) {
	if config == nil {
		config = &SessionConfig{}
	}
	config.applyDefault()
	session, err := client.NewSession()
	defer func() {
		if err != nil {
			session.Close()
		}
	}()
	if err != nil {
		return nil, err
	}
	for k, v := range config.EnvVariables {<commit_msg>added stdout drain before new command<commit_after>package ssh

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

const defaultShell = "/bin/bash"

const defautTimeoutMs = 5000

//MultiCommandSession represents a multi command session
//a new command are send vi stdin
type MultiCommandSession struct {
	session     *ssh.Session
	stdOutput   chan string
	stdError    chan string
	stdInput    io.WriteCloser
	ShellPrompt string
	running     int32
}

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

func (s *MultiCommandSession) init(shell string) (string, error) {
	reader, err := s.session.StdoutPipe()
	if err != nil {
		return "", err
	}
	go s.drain(reader, s.stdOutput)

	errReader, err := s.session.StderrPipe()
	if err != nil {
		return "", err
	}
	go s.drain(errReader, s.stdError)
	if shell == "" {
		shell = defaultShell
	}
	err = s.session.Start(shell)
	if err != nil {
		return "", err
	}
	return s.readResponse(defautTimeoutMs)
}

func (s *MultiCommandSession) drain(reader io.Reader, out chan string) {
	var written int64 = 0
	buf := make([]byte, 128*1024)
	for {
		writter := new(bytes.Buffer)
		if atomic.LoadInt32(&s.running) == 0 {
			return
		}

		bytesRead, readError := reader.Read(buf)
		if bytesRead > 0 {
			bytesWritten, writeError := writter.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(writter.Bytes())
		}
		if s.closeIfError(readError) {
			return
		}

	}
}

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

//Close closes the session with its resources
func (s *MultiCommandSession) Close() {
	atomic.StoreInt32(&s.running, 0)
	s.stdInput.Close()
	s.session.Close()

}

func (s *MultiCommandSession) readResponse(timeoutMs int, terminators ...string) (out string, err error) {
	if timeoutMs == 0 {
		timeoutMs = defautTimeoutMs
	}
	if len(terminators) == 0 {
		if s.ShellPrompt == "" {
			terminators = []string{s.ShellPrompt + "$"}
		} else {
			terminators = []string{"$ $"}
		}
	}
	var done int32
	defer atomic.StoreInt32(&done, 1)
	var errOut string
outer:
	for {
		select {

		case o := <-s.stdOutput:
			out += o
			if hasTerminator(out, terminators...) && len(s.stdOutput) == 0 {
				break outer
			}
		case e := <-s.stdError:
			errOut += e
			if hasTerminator(errOut, terminators...) && len(s.stdOutput) == 0 {
				break outer
			}

		case <-time.After(time.Duration(timeoutMs) * time.Millisecond):
			break outer
		}
	}
	if errOut != "" {
		err = errors.New(errOut)
	}
	if len(out) > 0 {
		index := strings.LastIndex(out, "\r\n"+s.ShellPrompt)
		if index > 0 {
			out = string(out[:index])
		}
	}
	return out, err
}

func (s *MultiCommandSession) drainStdout() {
	//read any outstanding output
	for ;; {
		out,  _:= s.readResponse(1, "")
		if len(out) == 0 {
			return
		}

	}
}

func (s *MultiCommandSession) Run(command string, timeoutMs int, terminators ...string) (string, error) {
	s.drainStdout()
	_, err := s.stdInput.Write([]byte(command + "\n"))
	if err != nil {
		return "", fmt.Errorf("Failed to execute command: %v, err: %v", command, err)
	}
	return s.readResponse(timeoutMs, terminators...)
}

func newMultiCommandSession(client *ssh.Client, config *SessionConfig) (result *MultiCommandSession, err error) {
	if config == nil {
		config = &SessionConfig{}
	}
	config.applyDefault()
	session, err := client.NewSession()
	defer func() {
		if err != nil {
			session.Close()
		}
	}()
	if err != nil {
		return nil, err
	}
	for k, v := range config.EnvVariables {
{'fraction_non_alphanumeric': 0.1425604690935744, 'fraction_numerical': 0.010750061079892499, 'mean_word_length': 3.4254054054054053, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1014994', 'n_tokens_mistral': 3216, 'n_tokens_neox': 2870, 'n_words': 731}