<commit_before>package main

import (
	"crypto/sha1"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

type mode int

const (
	stageFolder  = ".jazzstage"
	backupFolder = ".jazzbackup"

	STAGE mode = iota
	BACKUP
	NO_COPY
)

func statusDefaults() {
	fmt.Printf("gojazz status [options]\n")
	flag.PrintDefaults()
}

type status struct {
	Added    map[string]bool
	Modified map[string]bool
	Deleted  map[string]bool

	metaData *metaData

	sandboxPath string
	copyPath    string
}

func newStatus(sandboxPath string, m mode) *status {
	status := &status{}
	status.Added = make(map[string]bool)
	status.Modified = make(map[string]bool)
	status.Deleted = make(map[string]bool)

	status.sandboxPath = sandboxPath

	if m == STAGE {
		status.copyPath = filepath.Join(status.sandboxPath, stageFolder)
	} else if m == BACKUP {
		status.copyPath = filepath.Join(status.sandboxPath, backupFolder)
	}

	return status
}

func (status *status) unchanged() bool {
	return len(status.Added) == 0 && len(status.Modified) == 0 && len(status.Deleted) == 0
}

func (status *status) String() string {
	// TODO hit the server and find the current name of this workspace
	//result := status.metaData.workspaceName + "\n"

	result := ""

	if status.metaData.isstream {
		result = result + "Type: Stream\n"
	} else {
		result = result + "Type: Repository Workspace\n"
	}

	nochanges := true

	for k, _ := range status.Added {
		result = result + k + " (Added)\n"
		nochanges = false
	}

	for k, _ := range status.Modified {
		result = result + k + " (Modified)\n"
		nochanges = false
	}

	for k, _ := range status.Deleted {
		result = result + k + " (Deleted)\n"
		nochanges = false
	}

	if nochanges {
		result = result + "No local changes\n"
	}

	return result
}

func statusOp() {
	sandboxPath := flag.String("sandbox", "", "Location of the sandbox to load the files")
	flag.Usage = statusDefaults
	flag.Parse()

	if *sandboxPath == "" {
		path, err := os.Getwd()
		if err != nil {
			panic(err)
		}

		path = findSandbox(path)
		sandboxPath = &path
	}

	fmt.Printf("Status of %v...\n", *sandboxPath)
	status, err := scmStatus(*sandboxPath, NO_COPY)

	if err != nil {
		panic(err)
	}

	fmt.Printf("%v", status)
}

func scmStatus(sandboxPath string, m mode) (*status, error) {
	// Load up existing metadata and prepare fresh metadata
	oldMetaData := newMetaData()
	// If the load fails, it's not a problem, just empty
	err := oldMetaData.load(filepath.Join(sandboxPath, metadataFileName))

	if err != nil {
		return nil, simpleWarning("Not a sandbox")
	}

	status := newStatus(sandboxPath, m)
	status.metaData = oldMetaData

	// Delete any existing staging area
	if m == STAGE {
		err = os.RemoveAll(status.copyPath)
		if err != nil {
			panic(err)
		}
	}

	// Walk the current directory structure looking for changed items
	err = filepath.Walk(sandboxPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if path == sandboxPath {
			return nil
		}

		ignored, err := IsIgnored(path)
		if err != nil {
			return err
		}

		if ignored && info.IsDir() {
			return filepath.SkipDir
		} else if ignored {
			return nil
		}

		meta, ok := oldMetaData.get(path, sandboxPath)

		// Metadata doesn't exist for this file, so it must be added
		if !ok {
			status.fileAdded(path, sandboxPath)
			return nil
		}

		if !info.IsDir() {
			// Check the modified time
			if meta.LastModified != info.ModTime().Unix() {
				// Different sizes mean that the file has changed for sure
				if meta.Size != info.Size() {
					status.fileModified(meta, path, sandboxPath)
				} else {
					// Check the hashes
					file, err := os.Open(path)
					if err != nil {
						return err
					}

					hash := sha1.New()
					_, err = io.Copy(hash, file)

					if err != nil {
						return err
					}

					newHash := base64.StdEncoding.EncodeToString(hash.Sum(nil))

					if meta.Hash != newHash {
						status.fileModified(meta, path, sandboxPath)
					}
				}
			}
		}

		return nil
	})

	if err != nil {
		return nil, err
	}

	// Walk the metadata to find any items that don't exist
	for path, meta := range oldMetaData.pathMap {
		fullpath := filepath.Join(sandboxPath, path)
		_, err := os.Stat(fullpath)
		if err != nil {
			status.fileDeleted(meta, fullpath, sandboxPath)
		}
	}

	return status, nil
}

func (status *status) calcCopyPath(path string) string {
	if status.copyPath == "" {
		return ""
	}

	relpath, err := filepath.Rel(status.sandboxPath, path)

	if err != nil {
		panic(err)
	}

	return filepath.Join(status.copyPath, relpath)
}

func IsIgnored(path string) (bool, error) {
	// Should we ignore changes to this file?
	// Check for signs that it isn't source code
	//  -Really big file (>10MB)
	//  -File extension (e.g. .exe, .dll, .so)
	//  -Temporary files left by editors (e.g. *~, *.ext.swp)

	base := filepath.Base(path)

	// Skip the metadata, staging and backup directories
	if base == metadataFileName || strings.Contains(path, stageFolder) || strings.Contains(path, backupFolder) {
		return true, nil
	}

	// Skip bin directories
	if strings.HasSuffix(path, "/bin") || strings.Contains(path, "/bin/") {
		return true, nil
	}

	if strings.HasSuffix(base, ".exe") || strings.HasSuffix(base, ".dll") || strings.HasSuffix(base, ".so") {
		return true, nil
	}

	if strings.HasSuffix(base, "~") || strings.HasSuffix(base, ".ext.swp") {
		return true, nil
	}

	s, err := os.Stat(path)
	if err != nil {
		return false, err
	}

	if s.Size() > 10*1024*1024 {
		return true, nil
	}

	return false, nil
}

func (status *status) fileAdded(path string, sandboxPath string) {
	rel, err := filepath.Rel(sandboxPath, path)
	if err != nil {
		panic(err)
	}

	s, err := os.Stat(path)
	if err != nil {
		panic(err)
	}

	status.Added[rel] = true
	copyPath := status.calcCopyPath(path)

	if copyPath != "" {
		if s.IsDir() {
			os.MkdirAll(copyPath, 0700)
		} else {
			os.MkdirAll(filepath.Dir(copyPath), 0700)

			stagedFile, err := os.Create(copyPath)
			if err != nil {
				panic(err)
			}
			defer stagedFile.Close()
			origFile, err := os.Open(path)
			if err != nil {
				panic(err)
			}
			defer origFile.Close()

			_, err = io.Copy(stagedFile, origFile)
			if err != nil {
				panic(err)
			}
		}
	}
}

func (status *status) fileModified(meta metaObject, path string, sandboxPath string) {
	rel, err := filepath.Rel(sandboxPath, path)
	if err != nil {
		panic(err)
	}

	status.Modified[rel] = true
	copyPath := status.calcCopyPath(path)

	if copyPath != "" {
		s, err := os.Stat(path)

		if err != nil {
			panic(err)
		}

		if s.IsDir() {
			os.MkdirAll(copyPath, 0700)
		} else {
			os.MkdirAll(filepath.Dir(copyPath), 0700)

			stagedFile, err := os.Create(copyPath)
			if err != nil {
				panic(err)
			}
			defer stagedFile.Close()
			origFile, err := os.Open(path)
			if err != nil {
				panic(err)
			}
			defer origFile.Close()

			_, err = io.Copy(stagedFile, origFile)
			if err != nil {
				panic(err)
			}
		}
	}
}

func (status *status) fileDeleted(meta metaObject, path string, sandboxPath string) {
	rel, err := filepath.Rel(sandboxPath, path)
	if err != nil {
		panic(err)
	}

	status.Deleted[rel] = true
}
<commit_msg>Fix modification time check, improves automated test reliability. work item 35491<commit_after>package main

import (
	"crypto/sha1"
	"encoding/base64"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
)

type mode int

const (
	stageFolder  = ".jazzstage"
	backupFolder = ".jazzbackup"

	STAGE mode = iota
	BACKUP
	NO_COPY
)

func statusDefaults() {
	fmt.Printf("gojazz status [options]\n")
	flag.PrintDefaults()
}

type status struct {
	Added    map[string]bool
	Modified map[string]bool
	Deleted  map[string]bool

	metaData *metaData

	sandboxPath string
	copyPath    string
}

func newStatus(sandboxPath string, m mode) *status {
	status := &status{}
	status.Added = make(map[string]bool)
	status.Modified = make(map[string]bool)
	status.Deleted = make(map[string]bool)

	status.sandboxPath = sandboxPath

	if m == STAGE {
		status.copyPath = filepath.Join(status.sandboxPath, stageFolder)
	} else if m == BACKUP {
		status.copyPath = filepath.Join(status.sandboxPath, backupFolder)
	}

	return status
}

func (status *status) unchanged() bool {
	return len(status.Added) == 0 && len(status.Modified) == 0 && len(status.Deleted) == 0
}

func (status *status) String() string {
	// TODO hit the server and find the current name of this workspace
	//result := status.metaData.workspaceName + "\n"

	result := ""

	if status.metaData.isstream {
		result = result + "Type: Stream\n"
	} else {
		result = result + "Type: Repository Workspace\n"
	}

	nochanges := true

	for k, _ := range status.Added {
		result = result + k + " (Added)\n"
		nochanges = false
	}

	for k, _ := range status.Modified {
		result = result + k + " (Modified)\n"
		nochanges = false
	}

	for k, _ := range status.Deleted {
		result = result + k + " (Deleted)\n"
		nochanges = false
	}

	if nochanges {
		result = result + "No local changes\n"
	}

	return result
}

func statusOp() {
	sandboxPath := flag.String("sandbox", "", "Location of the sandbox to load the files")
	flag.Usage = statusDefaults
	flag.Parse()

	if *sandboxPath == "" {
		path, err := os.Getwd()
		if err != nil {
			panic(err)
		}

		path = findSandbox(path)
		sandboxPath = &path
	}

	fmt.Printf("Status of %v...\n", *sandboxPath)
	status, err := scmStatus(*sandboxPath, NO_COPY)

	if err != nil {
		panic(err)
	}

	fmt.Printf("%v", status)
}

func scmStatus(sandboxPath string, m mode) (*status, error) {
	// Load up existing metadata and prepare fresh metadata
	oldMetaData := newMetaData()
	// If the load fails, it's not a problem, just empty
	err := oldMetaData.load(filepath.Join(sandboxPath, metadataFileName))

	if err != nil {
		return nil, simpleWarning("Not a sandbox")
	}

	status := newStatus(sandboxPath, m)
	status.metaData = oldMetaData

	// Delete any existing staging area
	if m == STAGE {
		err = os.RemoveAll(status.copyPath)
		if err != nil {
			panic(err)
		}
	}

	// Walk the current directory structure looking for changed items
	err = filepath.Walk(sandboxPath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if path == sandboxPath {
			return nil
		}

		ignored, err := IsIgnored(path)
		if err != nil {
			return err
		}

		if ignored && info.IsDir() {
			return filepath.SkipDir
		} else if ignored {
			return nil
		}

		meta, ok := oldMetaData.get(path, sandboxPath)

		// Metadata doesn't exist for this file, so it must be added
		if !ok {
			status.fileAdded(path, sandboxPath)
			return nil
		}

		if !info.IsDir() {
			// The modified time is not a good enough check to see if the
			//  file is modified or not.
			//if meta.LastModified != info.ModTime().Unix() {
			// Different sizes mean that the file has changed for sure
			if meta.Size != info.Size() {
				status.fileModified(meta, path, sandboxPath)
			} else {
				// Check the hashes
				file, err := os.Open(path)
				if err != nil {
					return err
				}

				hash := sha1.New()
				_, err = io.Copy(hash, file)

				if err != nil {
					return err
				}

				newHash := base64.StdEncoding.EncodeToString(hash.Sum(nil))

				if meta.Hash != newHash {
					status.fileModified(meta, path, sandboxPath)
				}
			}
			//}
		}

		return nil
	})

	if err != nil {
		return nil, err
	}

	// Walk the metadata to find any items that don't exist
	for path, meta := range oldMetaData.pathMap {
		fullpath := filepath.Join(sandboxPath, path)
		_, err := os.Stat(fullpath)
		if err != nil {
			status.fileDeleted(meta, fullpath, sandboxPath)
		}
	}

	return status, nil
}

func (status *status) calcCopyPath(path string) string {
	if status.copyPath == "" {
		return ""
	}

	relpath, err := filepath.Rel(status.sandboxPath, path)

	if err != nil {
		panic(err)
	}

	return filepath.Join(status.copyPath, relpath)
}

func IsIgnored(path string) (bool, error) {
	// Should we ignore changes to this file?
	// Check for signs that it isn't source code
	//  -Really big file (>10MB)
	//  -File extension (e.g. .exe, .dll, .so)
	//  -Temporary files left by editors (e.g. *~, *.ext.swp)

	base := filepath.Base(path)

	// Skip the metadata, staging and backup directories
	if base == metadataFileName || strings.Contains(path, stageFolder) || strings.Contains(path, backupFolder) {
		return true, nil
	}

	// Skip bin directories
	if strings.HasSuffix(path, "/bin") || strings.Contains(path, "/bin/") {
		return true, nil
	}

	if strings.HasSuffix(base, ".exe") || strings.HasSuffix(base, ".dll") || strings.HasSuffix(base, ".so") {
		return true, nil
	}

	if strings.HasSuffix(base, "~") || strings.HasSuffix(base, ".ext.swp") {
		return true, nil
	}

	s, err := os.Stat(path)
	if err != nil {
		return false, err
	}

	if s.Size() > 10*1024*1024 {
		return true, nil
	}

	return false, nil
}

func (status *status) fileAdded(path string, sandboxPath string) {
	rel, err := filepath.Rel(sandboxPath, path)
	if err != nil {
		panic(err)
	}

	s, err := os.Stat(path)
	if err != nil {
		panic(err)
	}

	status.Added[rel] = true
	copyPath := status.calcCopyPath(path)

	if copyPath != "" {
		if s.IsDir() {
			os.MkdirAll(copyPath, 0700)
		} else {
			os.MkdirAll(filepath.Dir(copyPath), 0700)

			stagedFile, err := os.Create(copyPath)
			if err != nil {
				panic(err)
			}
			defer stagedFile.Close()
			origFile, err := os.Open(path)
			if err != nil {
				panic(err)
			}
			defer origFile.Close()

			_, err = io.Copy(stagedFile, origFile)
			if err != nil {
				panic(err)
			}
		}
	}
}

func (status *status) fileModified(meta metaObject, path string, sandboxPath string) {
	rel, err := filepath.Rel(sandboxPath, path)
	if err != nil {
		panic(err)
	}

	status.Modified[rel] = true
	copyPath := status.calcCopyPath(path)

	if copyPath != "" {
		s, err := os.Stat(path)

		if err != nil {
			panic(err)
		}

		if s.IsDir() {
			os.MkdirAll(copyPath, 0700)
		} else {
			os.MkdirAll(filepath.Dir(copyPath), 0700)

			stagedFile, err := os.Create(copyPath)
			if err != nil {
				panic(err)
			}
			defer stagedFile.Close()
			origFile, err := os.Open(path)
			if err != nil {
				panic(err)
			}
			defer origFile.Close()

			_, err = io.Copy(stagedFile, origFile)
			if err != nil {
				panic(err)
			}
		}
	}
}

func (status *status) fileDeleted(meta metaObject, path string, sandboxPath string) {
	rel, err := filepath.Rel(sandboxPath, path)
	if err != nil {
		panic(err)
	}

	status.Deleted[rel] = true
}
<|endoftext|>


<commit_before>package quote

import (
	"bytes"
	"encoding/xml"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/xtracdev/xavi/plugin"
	"github.com/xtracdev/xavi/plugin/timing"
	"github.com/xtracdev/xavisample/session"
	"golang.org/x/net/context"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"strings"
	"time"
)

func extractResource(uri string) (string, error) {
	parts := strings.Split(uri, "/")
	if len(parts) != 3 || parts[2] == "" {
		return "", fmt.Errorf("Expected URI format: /quote/<symbol>")
	}

	return parts[2], nil

}

func NewQuoteWrapper() plugin.Wrapper {
	return new(QuoteWrapper)
}

//For use in generating a variety of service names for use in exploring log management solutions,
//we'll randomly generate a service name.

var serviceNames = []string{"alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "india", "hotel"}

func generateServiceName() string {
	return serviceNames[rand.Intn(len(serviceNames))]
}

type QuoteWrapper struct{}

func (lw QuoteWrapper) Wrap(h plugin.ContextHandler) plugin.ContextHandler {
	return plugin.ContextHandlerFunc(func(c context.Context, w http.ResponseWriter, r *http.Request) {

		//Extract the timer from the service context
		end2endTimer := timing.TimerFromContext(c)
		if end2endTimer == nil {
			http.Error(w, "No timer in call context", http.StatusInternalServerError)
			return
		}

		//Set the top level name we want to use for recording timings, counts, etc.
		end2endTimer.Name = fmt.Sprintf("%s-quote", generateServiceName())

		contributor := end2endTimer.StartContributor("quote svc plugin")

		//Grab the symbol to quote from the uri
		resourceId, err := extractResource(r.RequestURI)
		if err != nil {
			contributor.End(err)
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte(err.Error()))
			return
		}

		if c != nil {
			sid, ok := c.Value(session.SessionKey).(int)
			if ok {
				log.Println("session:", sid, "symbol", resourceId)
			}
		}

		//Convert the method to POST for SOAP, and set the soap service
		//endpoint for the destination server
		r.Method = "POST"
		r.URL.Path = "/services/quote/getquote"

		//Form the SOAP payload
		payload := getQuoteRequestForSymbol(resourceId)
		payloadBytes, err := xml.Marshal(&payload)
		if err != nil {
			contributor.End(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		//Post the payload, and record the response
		r.Body = ioutil.NopCloser(bytes.NewReader(payloadBytes))
		rec := httptest.NewRecorder()

		h.ServeHTTPContext(c, rec, r)

		//Throw in a random service delay
		delay := rand.Intn(100) + 1
		time.Sleep(time.Duration(delay) * time.Millisecond)

		//Parse the recorded response to allow the quote price to be extracted
		var response ResponseEnvelope
		err = xml.Unmarshal(rec.Body.Bytes(), &response)
		if err != nil {
			contributor.End(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		//Return just the price to the caller
		w.Write([]byte(response.Body.GetLastTradePriceResponse.Price + "\n"))
		contributor.End(nil)
	})
}
<commit_msg>Updated sample to use better contributor name for quote service call<commit_after>package quote

import (
	"bytes"
	"encoding/xml"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/xtracdev/xavi/plugin"
	"github.com/xtracdev/xavi/plugin/timing"
	"github.com/xtracdev/xavisample/session"
	"golang.org/x/net/context"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"strings"
	"time"
)

func extractResource(uri string) (string, error) {
	parts := strings.Split(uri, "/")
	if len(parts) != 3 || parts[2] == "" {
		return "", fmt.Errorf("Expected URI format: /quote/<symbol>")
	}

	return parts[2], nil

}

func NewQuoteWrapper() plugin.Wrapper {
	return new(QuoteWrapper)
}

//For use in generating a variety of service names for use in exploring log management solutions,
//we'll randomly generate a service name.

var serviceNames = []string{"alpha", "bravo", "charlie", "delta", "echo", "foxtrot", "golf", "india", "hotel"}

func generateServiceName() string {
	return serviceNames[rand.Intn(len(serviceNames))]
}

type QuoteWrapper struct{}

func (lw QuoteWrapper) Wrap(h plugin.ContextHandler) plugin.ContextHandler {
	return plugin.ContextHandlerFunc(func(c context.Context, w http.ResponseWriter, r *http.Request) {

		//Extract the timer from the service context
		end2endTimer := timing.TimerFromContext(c)
		if end2endTimer == nil {
			http.Error(w, "No timer in call context", http.StatusInternalServerError)
			return
		}

		//Set the top level name we want to use for recording timings, counts, etc.
		end2endTimer.Name = fmt.Sprintf("%s-quote", generateServiceName())

		contributor := end2endTimer.StartContributor("QuoteSvc.GetTradePrice")

		//Grab the symbol to quote from the uri
		resourceId, err := extractResource(r.RequestURI)
		if err != nil {
			contributor.End(err)
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte(err.Error()))
			return
		}

		if c != nil {
			sid, ok := c.Value(session.SessionKey).(int)
			if ok {
				log.Println("session:", sid, "symbol", resourceId)
			}
		}

		//Convert the method to POST for SOAP, and set the soap service
		//endpoint for the destination server
		r.Method = "POST"
		r.URL.Path = "/services/quote/getquote"

		//Form the SOAP payload
		payload := getQuoteRequestForSymbol(resourceId)
		payloadBytes, err := xml.Marshal(&payload)
		if err != nil {
			contributor.End(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		//Post the payload, and record the response
		r.Body = ioutil.NopCloser(bytes.NewReader(payloadBytes))
		rec := httptest.NewRecorder()

		h.ServeHTTPContext(c, rec, r)

		//Throw in a random service delay
		delay := rand.Intn(100) + 1
		time.Sleep(time.Duration(delay) * time.Millisecond)

		//Parse the recorded response to allow the quote price to be extracted
		var response ResponseEnvelope
		err = xml.Unmarshal(rec.Body.Bytes(), &response)
		if err != nil {
			contributor.End(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		//Return just the price to the caller
		w.Write([]byte(response.Body.GetLastTradePriceResponse.Price + "\n"))
		contributor.End(nil)
	})
}
<|endoftext|>


<commit_before>package main

/* Heavily based off of https://github.com/dghubble/go-twitter/blob/master/examples/streaming.go */

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"

	"github.com/coreos/pkg/flagutil"
	"github.com/dghubble/go-twitter/twitter"
	"github.com/dghubble/oauth1"
)

func main() {
	flags := flag.NewFlagSet("user-auth", flag.ExitOnError)
	consumerKey := flags.String("consumer-key", "", "Twitter Consumer Key")
	consumerSecret := flags.String("consumer-secret", "", "Twitter Consumer Secret")
	accessToken := flags.String("access-token", "", "Twitter Access Token")
	accessSecret := flags.String("access-secret", "", "Twitter Access Secret")
	flags.Parse(os.Args[1:])
	flagutil.SetFlagsFromEnv(flags, "TWITTER")

	if *consumerKey == "" || *consumerSecret == "" || *accessToken == "" || *accessSecret == "" {
		log.Fatal("Consumer key/secret and Access token/secret required")
	}

	config := oauth1.NewConfig(*consumerKey, *consumerSecret)
	token := oauth1.NewToken(*accessToken, *accessSecret)
	// OAuth1 http.Client will automatically authorize Requests
	httpClient := config.Client(oauth1.NoContext, token)

	// Twitter Client
	client := twitter.NewClient(httpClient)

	// Convenience Demux demultiplexed stream messages
	// What to do with each type of tweet
	demux := twitter.NewSwitchDemux()

	demux.Tweet = func(tweet *twitter.Tweet) {
		fmt.Println("Pedestrian Down!")
		fmt.Println(tweet.Text)
		fmt.Printf("https://twitter.com/%s/status/%s\n", tweet.User.ScreenName, tweet.IDStr)

		//If not @yyctransport skip
		if tweet.User.ID != 1729579022 {
			return
		}

		// Determine suffix from number
		number := 0
		body, err := ioutil.ReadFile("/home/ubuntu/peddown/counter.txt")
		numberString := strings.TrimSpace(string(body))
		if err != nil {
			fmt.Println("counter file does not exist")
		} else {
			number, err = strconv.Atoi(numberString)
			if err != nil {
				fmt.Println("not a valid string")
				number = 1
			}
		}
		number += 1

		//Convert number to byte array and save the number
		numberString = strconv.Itoa(number)
		body = []byte(numberString)
		ioutil.WriteFile("/home/ubuntu/peddown/counter.txt", body, 0644)

		suffix := "th"
		switch number % 10 {
		case 1:
			suffix = "st"
		case 2:
			suffix = "nd"
		case 3:
			suffix = "rd"
		}

		tweetContent := fmt.Sprintf("%d%s tweeted pedestrian incident in 2016.\n#yycwalk #yyccc #ableg #visionzero\n@bmasonndp @shoffmanAB\n\nhttps://twitter.com/yyctransport/status/%s\n", number, suffix, tweet.IDStr)
		tweet, resp, err := client.Statuses.Update(tweetContent, nil)
		fmt.Println(resp)
		fmt.Println(err)
	}

	fmt.Println("Starting Stream...")

	// FILTER
	filterParams := &twitter.StreamFilterParams{
		Track:         []string{"ALERT ped,ALERT pedestrian"},
		StallWarnings: twitter.Bool(true),
	}
	stream, err := client.Streams.Filter(filterParams)
	if err != nil {
		log.Fatal(err)
	}

	// Receive messages until stopped or stream quits
	go demux.HandleChan(stream.Messages)

	// Wait for SIGINT and SIGTERM (HIT CTRL-C)
	ch := make(chan os.Signal)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
	log.Println(<-ch)

	fmt.Println("Stopping Stream...")
	stream.Stop()
}
<commit_msg>Ordinal math is hard apparently.<commit_after>package main

/* Heavily based off of https://github.com/dghubble/go-twitter/blob/master/examples/streaming.go */

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"

	"github.com/coreos/pkg/flagutil"
	"github.com/dghubble/go-twitter/twitter"
	"github.com/dghubble/oauth1"
)

func main() {
	flags := flag.NewFlagSet("user-auth", flag.ExitOnError)
	consumerKey := flags.String("consumer-key", "", "Twitter Consumer Key")
	consumerSecret := flags.String("consumer-secret", "", "Twitter Consumer Secret")
	accessToken := flags.String("access-token", "", "Twitter Access Token")
	accessSecret := flags.String("access-secret", "", "Twitter Access Secret")
	flags.Parse(os.Args[1:])
	flagutil.SetFlagsFromEnv(flags, "TWITTER")

	if *consumerKey == "" || *consumerSecret == "" || *accessToken == "" || *accessSecret == "" {
		log.Fatal("Consumer key/secret and Access token/secret required")
	}

	config := oauth1.NewConfig(*consumerKey, *consumerSecret)
	token := oauth1.NewToken(*accessToken, *accessSecret)
	// OAuth1 http.Client will automatically authorize Requests
	httpClient := config.Client(oauth1.NoContext, token)

	// Twitter Client
	client := twitter.NewClient(httpClient)

	// Convenience Demux demultiplexed stream messages
	// What to do with each type of tweet
	demux := twitter.NewSwitchDemux()

	demux.Tweet = func(tweet *twitter.Tweet) {
		fmt.Println("Pedestrian Down!")
		fmt.Println(tweet.Text)
		fmt.Printf("https://twitter.com/%s/status/%s\n", tweet.User.ScreenName, tweet.IDStr)

		//If not @yyctransport skip
		if tweet.User.ID != 1729579022 {
			return
		}

		// Determine suffix from number
		number := 0
		body, err := ioutil.ReadFile("/home/ubuntu/peddown/counter.txt")
		numberString := strings.TrimSpace(string(body))
		if err != nil {
			fmt.Println("counter file does not exist")
		} else {
			number, err = strconv.Atoi(numberString)
			if err != nil {
				fmt.Println("not a valid string")
				number = 1
			}
		}
		number += 1

		//Convert number to byte array and save the number
		numberString = strconv.Itoa(number)
		body = []byte(numberString)
		ioutil.WriteFile("/home/ubuntu/peddown/counter.txt", body, 0644)

		suffix := "th"
		switch number % 10 {
		case 1:
			if (number % 100) != 11 {
				suffix = "st"
			}
		case 2:
			if (number % 100) != 12 {
				suffix = "nd"
			}
		case 3:
			if (number % 100) != 13 {
				suffix = "rd"
			}
		}

		tweetContent := fmt.Sprintf("%d%s tweeted pedestrian incident in 2016.\n#yycwalk #yyccc #ableg #visionzero\n@bmasonndp @shoffmanAB\n\nhttps://twitter.com/yyctransport/status/%s\n", number, suffix, tweet.IDStr)
		tweet, resp, err := client.Statuses.Update(tweetContent, nil)
		fmt.Println(resp)
		fmt.Println(err)
	}

	fmt.Println("Starting Stream...")

	// FILTER
	filterParams := &twitter.StreamFilterParams{
		Track:         []string{"ALERT ped,ALERT pedestrian"},
		StallWarnings: twitter.Bool(true),
	}
	stream, err := client.Streams.Filter(filterParams)
	if err != nil {
		log.Fatal(err)
	}

	// Receive messages until stopped or stream quits
	go demux.HandleChan(stream.Messages)

	// Wait for SIGINT and SIGTERM (HIT CTRL-C)
	ch := make(chan os.Signal)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
	log.Println(<-ch)

	fmt.Println("Stopping Stream...")
	stream.Stop()
}
<|endoftext|>


<commit_before>package campfire

import (
    "github.com/brettbuddin/httpie"
    "net/url"
    "fmt"
    "encoding/json"
    "log"
)

type Stream struct {
    base     *httpie.Stream
    room     *Room
    messages chan *Message
}

// NewStream returns a Stream
func NewStream(room *Room, messages chan *Message) *Stream {
    return &Stream{
        room: room,
        messages: messages,
    }
}

// Connect starts the stream
func (s *Stream) Connect() {
    url := &url.URL{
        Scheme: "https",
        Host: "streaming.campfirenow.com",
        Path: fmt.Sprintf("/room/%d/live.json", s.room.Id),
    }

    s.base = httpie.NewStream(
        httpie.Get{url},
        httpie.BasicAuth{s.room.conn.token, "X"},
        httpie.CarriageReturn,
    )

    go s.base.Connect()

    log.Println("Streaming")
    for {
        select {
        case data := <-s.base.Data():
            var m Message
            err := json.Unmarshal(data, &m)

            if err != nil {
                continue
            }

            m.conn = s.room.conn
            s.messages <- &m
        }
    }
}

// Disconnect stops the stream
func (s *Stream) Disconnect() {
    s.base.Disconnect()
}
<commit_msg>Don't spit this out<commit_after>package campfire

import (
    "github.com/brettbuddin/httpie"
    "net/url"
    "fmt"
    "encoding/json"
)

type Stream struct {
    base     *httpie.Stream
    room     *Room
    messages chan *Message
}

// NewStream returns a Stream
func NewStream(room *Room, messages chan *Message) *Stream {
    return &Stream{
        room: room,
        messages: messages,
    }
}

// Connect starts the stream
func (s *Stream) Connect() {
    url := &url.URL{
        Scheme: "https",
        Host: "streaming.campfirenow.com",
        Path: fmt.Sprintf("/room/%d/live.json", s.room.Id),
    }

    s.base = httpie.NewStream(
        httpie.Get{url},
        httpie.BasicAuth{s.room.conn.token, "X"},
        httpie.CarriageReturn,
    )

    go s.base.Connect()

    for {
        select {
        case data := <-s.base.Data():
            var m Message
            err := json.Unmarshal(data, &m)

            if err != nil {
                continue
            }

            m.conn = s.room.conn
            s.messages <- &m
        }
    }
}

// Disconnect stops the stream
func (s *Stream) Disconnect() {
    s.base.Disconnect()
}
<|endoftext|>


<commit_before>package yamux

import (
	"bytes"
	"io"
	"sync"
	"sync/atomic"
	"time"
)

type streamState int

const (
	streamInit streamState = iota
	streamSYNSent
	streamSYNReceived
	streamEstablished
	streamLocalClose
	streamRemoteClose
	streamClosed
	streamReset
)

// Stream is used to represent a logical stream
// within a session.
type Stream struct {
	recvWindow uint32
	sendWindow uint32

	id      uint32
	session *Session

	state     streamState
	stateLock sync.Mutex

	recvBuf  bytes.Buffer
	recvLock sync.Mutex

	controlHdr     header
	controlErr     chan error
	controlHdrLock sync.Mutex

	sendHdr  header
	sendErr  chan error
	sendLock sync.Mutex

	recvNotifyCh chan struct{}
	sendNotifyCh chan struct{}

	readDeadline  time.Time
	writeDeadline time.Time
}

// newStream is used to construct a new stream within
// a given session for an ID
func newStream(session *Session, id uint32, state streamState) *Stream {
	s := &Stream{
		id:           id,
		session:      session,
		state:        state,
		controlHdr:   header(make([]byte, headerSize)),
		controlErr:   make(chan error, 1),
		sendHdr:      header(make([]byte, headerSize)),
		sendErr:      make(chan error, 1),
		recvWindow:   initialStreamWindow,
		sendWindow:   initialStreamWindow,
		recvNotifyCh: make(chan struct{}, 1),
		sendNotifyCh: make(chan struct{}, 1),
	}
	return s
}

// Session returns the associated stream session
func (s *Stream) Session() *Session {
	return s.session
}

// StreamID returns the ID of this stream
func (s *Stream) StreamID() uint32 {
	return s.id
}

// Read is used to read from the stream
func (s *Stream) Read(b []byte) (n int, err error) {
	defer asyncNotify(s.recvNotifyCh)
START:
	s.stateLock.Lock()
	switch s.state {
	case streamRemoteClose:
		fallthrough
	case streamClosed:
		if s.recvBuf.Len() == 0 {
			s.stateLock.Unlock()
			return 0, io.EOF
		}
	case streamReset:
		s.stateLock.Unlock()
		return 0, ErrConnectionReset
	}
	s.stateLock.Unlock()

	// If there is no data available, block
	s.recvLock.Lock()
	if s.recvBuf.Len() == 0 {
		s.recvLock.Unlock()
		goto WAIT
	}

	// Read any bytes
	n, _ = s.recvBuf.Read(b)
	s.recvLock.Unlock()

	// Send a window update potentially
	err = s.sendWindowUpdate()
	return n, err

WAIT:
	var timeout <-chan time.Time
	if !s.readDeadline.IsZero() {
		delay := s.readDeadline.Sub(time.Now())
		timeout = time.After(delay)
	}
	select {
	case <-s.recvNotifyCh:
		goto START
	case <-timeout:
		return 0, ErrTimeout
	}
}

// Write is used to write to the stream
func (s *Stream) Write(b []byte) (n int, err error) {
	s.sendLock.Lock()
	defer s.sendLock.Unlock()
	total := 0
	for total < len(b) {
		n, err := s.write(b[total:])
		total += n
		if err != nil {
			return total, err
		}
	}
	return total, nil
}

// write is used to write to the stream, may return on
// a short write.
func (s *Stream) write(b []byte) (n int, err error) {
	var flags uint16
	var max uint32
	var body io.Reader
START:
	s.stateLock.Lock()
	switch s.state {
	case streamLocalClose:
		fallthrough
	case streamClosed:
		s.stateLock.Unlock()
		return 0, ErrStreamClosed
	case streamReset:
		s.stateLock.Unlock()
		return 0, ErrConnectionReset
	}
	s.stateLock.Unlock()

	// If there is no data available, block
	window := atomic.LoadUint32(&s.sendWindow)
	if window == 0 {
		goto WAIT
	}

	// Determine the flags if any
	flags = s.sendFlags()

	// Send up to our send window
	max = min(window, uint32(len(b)))
	body = bytes.NewReader(b[:max])

	// Send the header
	s.sendHdr.encode(typeData, flags, s.id, max)
	if err := s.session.waitForSendErr(s.sendHdr, body, s.sendErr); err != nil {
		return 0, err
	}

	// Reduce our send window
	atomic.AddUint32(&s.sendWindow, ^uint32(max-1))

	// Unlock
	return int(max), err

WAIT:
	var timeout <-chan time.Time
	if !s.writeDeadline.IsZero() {
		delay := s.writeDeadline.Sub(time.Now())
		timeout = time.After(delay)
	}
	select {
	case <-s.sendNotifyCh:
		goto START
	case <-timeout:
		return 0, ErrTimeout
	}
	return 0, nil
}

// sendFlags determines any flags that are appropriate
// based on the current stream state
func (s *Stream) sendFlags() uint16 {
	s.stateLock.Lock()
	defer s.stateLock.Unlock()
	var flags uint16
	switch s.state {
	case streamInit:
		flags |= flagSYN
		s.state = streamSYNSent
	case streamSYNReceived:
		flags |= flagACK
		s.state = streamEstablished
	}
	return flags
}

// sendWindowUpdate potentially sends a window update enabling
// further writes to take place. Must be invoked with the lock.
func (s *Stream) sendWindowUpdate() error {
	s.controlHdrLock.Lock()
	defer s.controlHdrLock.Unlock()

	// Determine the delta update
	max := s.session.config.MaxStreamWindowSize
	delta := max - atomic.LoadUint32(&s.recvWindow)

	// Determine the flags if any
	flags := s.sendFlags()

	// Check if we can omit the update
	if delta < (max/2) && flags == 0 {
		return nil
	}

	// Update our window
	atomic.AddUint32(&s.recvWindow, delta)

	// Send the header
	s.controlHdr.encode(typeWindowUpdate, flags, s.id, delta)
	if err := s.session.waitForSendErr(s.controlHdr, nil, s.controlErr); err != nil {
		return err
	}
	return nil
}

// sendClose is used to send a FIN
func (s *Stream) sendClose() error {
	s.controlHdrLock.Lock()
	defer s.controlHdrLock.Unlock()

	flags := s.sendFlags()
	flags |= flagFIN
	s.controlHdr.encode(typeWindowUpdate, flags, s.id, 0)
	if err := s.session.waitForSendErr(s.controlHdr, nil, s.controlErr); err != nil {
		return err
	}
	return nil
}

// Close is used to close the stream
func (s *Stream) Close() error {
	s.stateLock.Lock()
	switch s.state {
	// Opened means we need to signal a close
	case streamSYNSent:
		fallthrough
	case streamSYNReceived:
		fallthrough
	case streamEstablished:
		s.state = streamLocalClose
		goto SEND_CLOSE

	case streamLocalClose:
	case streamRemoteClose:
		s.state = streamClosed
		s.session.closeStream(s.id)
		goto SEND_CLOSE

	case streamClosed:
	case streamReset:
	default:
		panic("unhandled state")
	}
	s.stateLock.Unlock()
	return nil
SEND_CLOSE:
	s.stateLock.Unlock()
	s.sendClose()
	s.notifyWaiting()
	return nil
}

// forceClose is used for when the session is exiting
func (s *Stream) forceClose() {
	s.stateLock.Lock()
	s.state = streamClosed
	s.stateLock.Unlock()
	s.notifyWaiting()
}

// processFlags is used to update the state of the stream
// based on set flags, if any. Lock must be held
func (s *Stream) processFlags(flags uint16) error {
	s.stateLock.Lock()
	defer s.stateLock.Unlock()
	if flags&flagACK == flagACK {
		if s.state == streamSYNSent {
			s.state = streamEstablished
		}

	} else if flags&flagFIN == flagFIN {
		switch s.state {
		case streamSYNSent:
			fallthrough
		case streamSYNReceived:
			fallthrough
		case streamEstablished:
			s.state = streamRemoteClose
			s.notifyWaiting()
		case streamLocalClose:
			s.state = streamClosed
			s.session.closeStream(s.id)
			s.notifyWaiting()
		default:
			s.session.logger.Printf("[ERR] yamux: unexpected FIN flag in state %d", s.state)
			return ErrUnexpectedFlag
		}
	} else if flags&flagRST == flagRST {
		s.state = streamReset
		s.session.closeStream(s.id)
		s.notifyWaiting()
	}
	return nil
}

// notifyWaiting notifies all the waiting channels
func (s *Stream) notifyWaiting() {
	asyncNotify(s.recvNotifyCh)
	asyncNotify(s.sendNotifyCh)
}

// incrSendWindow updates the size of our send window
func (s *Stream) incrSendWindow(hdr header, flags uint16) error {
	if err := s.processFlags(flags); err != nil {
		return err
	}

	// Increase window, unblock a sender
	atomic.AddUint32(&s.sendWindow, hdr.Length())
	asyncNotify(s.sendNotifyCh)
	return nil
}

// readData is used to handle a data frame
func (s *Stream) readData(hdr header, flags uint16, conn io.Reader) error {
	if err := s.processFlags(flags); err != nil {
		return err
	}

	// Check that our recv window is not exceeded
	length := hdr.Length()
	if length == 0 {
		return nil
	}
	if remain := atomic.LoadUint32(&s.recvWindow); length > remain {
		s.session.logger.Printf("[ERR] yamux: receive window exceeded (stream: %d, remain: %d, recv: %d)", s.id, remain, length)
		return ErrRecvWindowExceeded
	}

	// Wrap in a limited reader
	conn = &io.LimitedReader{R: conn, N: int64(length)}

	// Copy into buffer
	s.recvLock.Lock()
	if _, err := io.Copy(&s.recvBuf, conn); err != nil {
		s.session.logger.Printf("[ERR] yamux: Failed to read stream data: %v", err)
		s.recvLock.Unlock()
		return err
	}

	// Decrement the receive window
	atomic.AddUint32(&s.recvWindow, ^uint32(length-1))
	s.recvLock.Unlock()

	// Unblock any readers
	asyncNotify(s.recvNotifyCh)
	return nil
}

// SetDeadline sets the read and write deadlines
func (s *Stream) SetDeadline(t time.Time) error {
	if err := s.SetReadDeadline(t); err != nil {
		return err
	}
	if err := s.SetWriteDeadline(t); err != nil {
		return err
	}
	return nil
}

// SetReadDeadline sets the deadline for future Read calls.
func (s *Stream) SetReadDeadline(t time.Time) error {
	s.readDeadline = t
	return nil
}

// SetWriteDeadline sets the deadline for future Write calls
func (s *Stream) SetWriteDeadline(t time.Time) error {
	s.writeDeadline = t
	return nil
}
<commit_msg>Prevent Read on a closed stream<commit_after>package yamux

import (
	"bytes"
	"io"
	"sync"
	"sync/atomic"
	"time"
)

type streamState int

const (
	streamInit streamState = iota
	streamSYNSent
	streamSYNReceived
	streamEstablished
	streamLocalClose
	streamRemoteClose
	streamClosed
	streamReset
)

// Stream is used to represent a logical stream
// within a session.
type Stream struct {
	recvWindow uint32
	sendWindow uint32

	id      uint32
	session *Session

	state     streamState
	stateLock sync.Mutex

	recvBuf  bytes.Buffer
	recvLock sync.Mutex

	controlHdr     header
	controlErr     chan error
	controlHdrLock sync.Mutex

	sendHdr  header
	sendErr  chan error
	sendLock sync.Mutex

	recvNotifyCh chan struct{}
	sendNotifyCh chan struct{}

	readDeadline  time.Time
	writeDeadline time.Time
}

// newStream is used to construct a new stream within
// a given session for an ID
func newStream(session *Session, id uint32, state streamState) *Stream {
	s := &Stream{
		id:           id,
		session:      session,
		state:        state,
		controlHdr:   header(make([]byte, headerSize)),
		controlErr:   make(chan error, 1),
		sendHdr:      header(make([]byte, headerSize)),
		sendErr:      make(chan error, 1),
		recvWindow:   initialStreamWindow,
		sendWindow:   initialStreamWindow,
		recvNotifyCh: make(chan struct{}, 1),
		sendNotifyCh: make(chan struct{}, 1),
	}
	return s
}

// Session returns the associated stream session
func (s *Stream) Session() *Session {
	return s.session
}

// StreamID returns the ID of this stream
func (s *Stream) StreamID() uint32 {
	return s.id
}

// Read is used to read from the stream
func (s *Stream) Read(b []byte) (n int, err error) {
	defer asyncNotify(s.recvNotifyCh)
START:
	s.stateLock.Lock()
	switch s.state {
	case streamLocalClose:
		fallthrough
	case streamRemoteClose:
		fallthrough
	case streamClosed:
		if s.recvBuf.Len() == 0 {
			s.stateLock.Unlock()
			return 0, io.EOF
		}
	case streamReset:
		s.stateLock.Unlock()
		return 0, ErrConnectionReset
	}
	s.stateLock.Unlock()

	// If there is no data available, block
	s.recvLock.Lock()
	if s.recvBuf.Len() == 0 {
		s.recvLock.Unlock()
		goto WAIT
	}

	// Read any bytes
	n, _ = s.recvBuf.Read(b)
	s.recvLock.Unlock()

	// Send a window update potentially
	err = s.sendWindowUpdate()
	return n, err

WAIT:
	var timeout <-chan time.Time
	if !s.readDeadline.IsZero() {
		delay := s.readDeadline.Sub(time.Now())
		timeout = time.After(delay)
	}
	select {
	case <-s.recvNotifyCh:
		goto START
	case <-timeout:
		return 0, ErrTimeout
	}
}

// Write is used to write to the stream
func (s *Stream) Write(b []byte) (n int, err error) {
	s.sendLock.Lock()
	defer s.sendLock.Unlock()
	total := 0
	for total < len(b) {
		n, err := s.write(b[total:])
		total += n
		if err != nil {
			return total, err
		}
	}
	return total, nil
}

// write is used to write to the stream, may return on
// a short write.
func (s *Stream) write(b []byte) (n int, err error) {
	var flags uint16
	var max uint32
	var body io.Reader
START:
	s.stateLock.Lock()
	switch s.state {
	case streamLocalClose:
		fallthrough
	case streamClosed:
		s.stateLock.Unlock()
		return 0, ErrStreamClosed
	case streamReset:
		s.stateLock.Unlock()
		return 0, ErrConnectionReset
	}
	s.stateLock.Unlock()

	// If there is no data available, block
	window := atomic.LoadUint32(&s.sendWindow)
	if window == 0 {
		goto WAIT
	}

	// Determine the flags if any
	flags = s.sendFlags()

	// Send up to our send window
	max = min(window, uint32(len(b)))
	body = bytes.NewReader(b[:max])

	// Send the header
	s.sendHdr.encode(typeData, flags, s.id, max)
	if err := s.session.waitForSendErr(s.sendHdr, body, s.sendErr); err != nil {
		return 0, err
	}

	// Reduce our send window
	atomic.AddUint32(&s.sendWindow, ^uint32(max-1))

	// Unlock
	return int(max), err

WAIT:
	var timeout <-chan time.Time
	if !s.writeDeadline.IsZero() {
		delay := s.writeDeadline.Sub(time.Now())
		timeout = time.After(delay)
	}
	select {
	case <-s.sendNotifyCh:
		goto START
	case <-timeout:
		return 0, ErrTimeout
	}
	return 0, nil
}

// sendFlags determines any flags that are appropriate
// based on the current stream state
func (s *Stream) sendFlags() uint16 {
	s.stateLock.Lock()
	defer s.stateLock.Unlock()
	var flags uint16
	switch s.state {
	case streamInit:
		flags |= flagSYN
		s.state = streamSYNSent
	case streamSYNReceived:
		flags |= flagACK
		s.state = streamEstablished
	}
	return flags
}

// sendWindowUpdate potentially sends a window update enabling
// further writes to take place. Must be invoked with the lock.
func (s *Stream) sendWindowUpdate() error {
	s.controlHdrLock.Lock()
	defer s.controlHdrLock.Unlock()

	// Determine the delta update
	max := s.session.config.MaxStreamWindowSize
	delta := max - atomic.LoadUint32(&s.recvWindow)

	// Determine the flags if any
	flags := s.sendFlags()

	// Check if we can omit the update
	if delta < (max/2) && flags == 0 {
		return nil
	}

	// Update our window
	atomic.AddUint32(&s.recvWindow, delta)

	// Send the header
	s.controlHdr.encode(typeWindowUpdate, flags, s.id, delta)
	if err := s.session.waitForSendErr(s.controlHdr, nil, s.controlErr); err != nil {
		return err
	}
	return nil
}

// sendClose is used to send a FIN
func (s *Stream) sendClose() error {
	s.controlHdrLock.Lock()
	defer s.controlHdrLock.Unlock()

	flags := s.sendFlags()
	flags |= flagFIN
	s.controlHdr.encode(typeWindowUpdate, flags, s.id, 0)
	if err := s.session.waitForSendErr(s.controlHdr, nil, s.controlErr); err != nil {
		return err
	}
	return nil
}

// Close is used to close the stream
func (s *Stream) Close() error {
	s.stateLock.Lock()
	switch s.state {
	// Opened means we need to signal a close
	case streamSYNSent:
		fallthrough
	case streamSYNReceived:
		fallthrough
	case streamEstablished:
		s.state = streamLocalClose
		goto SEND_CLOSE

	case streamLocalClose:
	case streamRemoteClose:
		s.state = streamClosed
		s.session.closeStream(s.id)
		goto SEND_CLOSE

	case streamClosed:
	case streamReset:
	default:
		panic("unhandled state")
	}
	s.stateLock.Unlock()
	return nil
SEND_CLOSE:
	s.stateLock.Unlock()
	s.sendClose()
	s.notifyWaiting()
	return nil
}

// forceClose is used for when the session is exiting
func (s *Stream) forceClose() {
	s.stateLock.Lock()
	s.state = streamClosed
	s.stateLock.Unlock()
	s.notifyWaiting()
}

// processFlags is used to update the state of the stream
// based on set flags, if any. Lock must be held
func (s *Stream) processFlags(flags uint16) error {
	s.stateLock.Lock()
	defer s.stateLock.Unlock()
	if flags&flagACK == flagACK {
		if s.state == streamSYNSent {
			s.state = streamEstablished
		}

	} else if flags&flagFIN == flagFIN {
		switch s.state {
		case streamSYNSent:
			fallthrough
		case streamSYNReceived:
			fallthrough
		case streamEstablished:
			s.state = streamRemoteClose
			s.notifyWaiting()
		case streamLocalClose:
			s.state = streamClosed
			s.session.closeStream(s.id)
			s.notifyWaiting()
		default:
			s.session.logger.Printf("[ERR] yamux: unexpected FIN flag in state %d", s.state)
			return ErrUnexpectedFlag
		}
	} else if flags&flagRST == flagRST {
		s.state = streamReset
		s.session.closeStream(s.id)
		s.notifyWaiting()
	}
	return nil
}

// notifyWaiting notifies all the waiting channels
func (s *Stream) notifyWaiting() {
	asyncNotify(s.recvNotifyCh)
	asyncNotify(s.sendNotifyCh)
}

// incrSendWindow updates the size of our send window
func (s *Stream) incrSendWindow(hdr header, flags uint16) error {
	if err := s.processFlags(flags); err != nil {
		return err
	}

	// Increase window, unblock a sender
	atomic.AddUint32(&s.sendWindow, hdr.Length())
	asyncNotify(s.sendNotifyCh)
	return nil
}

// readData is used to handle a data frame
func (s *Stream) readData(hdr header, flags uint16, conn io.Reader) error {
	if err := s.processFlags(flags); err != nil {
		return err
	}

	// Check that our recv window is not exceeded
	length := hdr.Length()
	if length == 0 {
		return nil
	}
	if remain := atomic.LoadUint32(&s.recvWindow); length > remain {
		s.session.logger.Printf("[ERR] yamux: receive window exceeded (stream: %d, remain: %d, recv: %d)", s.id, remain, length)
		return ErrRecvWindowExceeded
	}

	// Wrap in a limited reader
	conn = &io.LimitedReader{R: conn, N: int64(length)}

	// Copy into buffer
	s.recvLock.Lock()
	if _, err := io.Copy(&s.recvBuf, conn); err != nil {
		s.session.logger.Printf("[ERR] yamux: Failed to read stream data: %v", err)
		s.recvLock.Unlock()
		return err
	}

	// Decrement the receive window
	atomic.AddUint32(&s.recvWindow, ^uint32(length-1))
	s.recvLock.Unlock()

	// Unblock any readers
	asyncNotify(s.recvNotifyCh)
	return nil
}

// SetDeadline sets the read and write deadlines
func (s *Stream) SetDeadline(t time.Time) error {
	if err := s.SetReadDeadline(t); err != nil {
		return err
	}
	if err := s.SetWriteDeadline(t); err != nil {
		return err
	}
	return nil
}

// SetReadDeadline sets the deadline for future Read calls.
func (s *Stream) SetReadDeadline(t time.Time) error {
	s.readDeadline = t
	return nil
}

// SetWriteDeadline sets the deadline for future Write calls
func (s *Stream) SetWriteDeadline(t time.Time) error {
	s.writeDeadline = t
	return nil
}
<|endoftext|>


<commit_before>package yamux

import (
	"bytes"
	"io"
	"sync"
	"sync/atomic"
	"time"
)

type streamState int

const (
	streamInit streamState = iota
	streamSYNSent
	streamSYNReceived
	streamEstablished
	streamLocalClose
	streamRemoteClose
	streamClosed
	streamReset
)

// Stream is used to represent a logical stream
// within a session.
type Stream struct {
	recvWindow uint32
	sendWindow uint32

	id      uint32
	session *Session

	state     streamState
	stateLock sync.Mutex

	recvBuf  *bytes.Buffer
	recvLock sync.Mutex

	controlHdr     header
	controlErr     chan error
	controlHdrLock sync.Mutex

	sendHdr  header
	sendErr  chan error
	sendLock sync.Mutex

	recvNotifyCh chan struct{}
	sendNotifyCh chan struct{}

	readDeadline  time.Time
	writeDeadline time.Time
}

// newStream is used to construct a new stream within
// a given session for an ID
func newStream(session *Session, id uint32, state streamState) *Stream {
	s := &Stream{
		id:           id,
		session:      session,
		state:        state,
		controlHdr:   header(make([]byte, headerSize)),
		controlErr:   make(chan error, 1),
		sendHdr:      header(make([]byte, headerSize)),
		sendErr:      make(chan error, 1),
		recvWindow:   initialStreamWindow,
		sendWindow:   initialStreamWindow,
		recvNotifyCh: make(chan struct{}, 1),
		sendNotifyCh: make(chan struct{}, 1),
	}
	return s
}

// Session returns the associated stream session
func (s *Stream) Session() *Session {
	return s.session
}

// StreamID returns the ID of this stream
func (s *Stream) StreamID() uint32 {
	return s.id
}

// Read is used to read from the stream
func (s *Stream) Read(b []byte) (n int, err error) {
	defer asyncNotify(s.recvNotifyCh)
START:
	s.stateLock.Lock()
	switch s.state {
	case streamLocalClose:
		fallthrough
	case streamRemoteClose:
		fallthrough
	case streamClosed:
		if s.recvBuf == nil || s.recvBuf.Len() == 0 {
			s.stateLock.Unlock()
			return 0, io.EOF
		}
	case streamReset:
		s.stateLock.Unlock()
		return 0, ErrConnectionReset
	}
	s.stateLock.Unlock()

	// If there is no data available, block
	s.recvLock.Lock()
	if s.recvBuf == nil || s.recvBuf.Len() == 0 {
		s.recvLock.Unlock()
		goto WAIT
	}

	// Read any bytes
	n, _ = s.recvBuf.Read(b)
	s.recvLock.Unlock()

	// Send a window update potentially
	err = s.sendWindowUpdate()
	return n, err

WAIT:
	var timeout <-chan time.Time
	if !s.readDeadline.IsZero() {
		delay := s.readDeadline.Sub(time.Now())
		timeout = time.After(delay)
	}
	select {
	case <-s.recvNotifyCh:
		goto START
	case <-timeout:
		return 0, ErrTimeout
	}
}

// Write is used to write to the stream
func (s *Stream) Write(b []byte) (n int, err error) {
	s.sendLock.Lock()
	defer s.sendLock.Unlock()
	total := 0
	for total < len(b) {
		n, err := s.write(b[total:])
		total += n
		if err != nil {
			return total, err
		}
	}
	return total, nil
}

// write is used to write to the stream, may return on
// a short write.
func (s *Stream) write(b []byte) (n int, err error) {
	var flags uint16
	var max uint32
	var body io.Reader
START:
	s.stateLock.Lock()
	switch s.state {
	case streamLocalClose:
		fallthrough
	case streamClosed:
		s.stateLock.Unlock()
		return 0, ErrStreamClosed
	case streamReset:
		s.stateLock.Unlock()
		return 0, ErrConnectionReset
	}
	s.stateLock.Unlock()

	// If there is no data available, block
	window := atomic.LoadUint32(&s.sendWindow)
	if window == 0 {
		goto WAIT
	}

	// Determine the flags if any
	flags = s.sendFlags()

	// Send up to our send window
	max = min(window, uint32(len(b)))
	body = bytes.NewReader(b[:max])

	// Send the header
	s.sendHdr.encode(typeData, flags, s.id, max)
	if err := s.session.waitForSendErr(s.sendHdr, body, s.sendErr); err != nil {
		return 0, err
	}

	// Reduce our send window
	atomic.AddUint32(&s.sendWindow, ^uint32(max-1))

	// Unlock
	return int(max), err

WAIT:
	var timeout <-chan time.Time
	if !s.writeDeadline.IsZero() {
		delay := s.writeDeadline.Sub(time.Now())
		timeout = time.After(delay)
	}
	select {
	case <-s.sendNotifyCh:
		goto START
	case <-timeout:
		return 0, ErrTimeout
	}
	return 0, nil
}

// sendFlags determines any flags that are appropriate
// based on the current stream state
func (s *Stream) sendFlags() uint16 {
	s.stateLock.Lock()
	defer s.stateLock.Unlock()
	var flags uint16
	switch s.state {
	case streamInit:
		flags |= flagSYN
		s.state = streamSYNSent
	case streamSYNReceived:
		flags |= flagACK
		s.state = streamEstablished
	}
	return flags
}

// sendWindowUpdate potentially sends a window update enabling
// further writes to take place. Must be invoked with the lock.
func (s *Stream) sendWindowUpdate() error {
	s.controlHdrLock.Lock()
	defer s.controlHdrLock.Unlock()

	// Determine the delta update
	max := s.session.config.MaxStreamWindowSize
	delta := max - atomic.LoadUint32(&s.recvWindow)

	// Determine the flags if any
	flags := s.sendFlags()

	// Check if we can omit the update
	if delta < (max/2) && flags == 0 {
		return nil
	}

	// Update our window
	atomic.AddUint32(&s.recvWindow, delta)

	// Send the header
	s.controlHdr.encode(typeWindowUpdate, flags, s.id, delta)
	if err := s.session.waitForSendErr(s.controlHdr, nil, s.controlErr); err != nil {
		return err
	}
	return nil
}

// sendClose is used to send a FIN
func (s *Stream) sendClose() error {
	s.controlHdrLock.Lock()
	defer s.controlHdrLock.Unlock()

	flags := s.sendFlags()
	flags |= flagFIN
	s.controlHdr.encode(typeWindowUpdate, flags, s.id, 0)
	if err := s.session.waitForSendErr(s.controlHdr, nil, s.controlErr); err != nil {
		return err
	}
	return nil
}

// Close is used to close the stream
func (s *Stream) Close() error {
	closeStream := false
	s.stateLock.Lock()
	switch s.state {
	// Opened means we need to signal a close
	case streamSYNSent:
		fallthrough
	case streamSYNReceived:
		fallthrough
	case streamEstablished:
		s.state = streamLocalClose
		goto SEND_CLOSE

	case streamLocalClose:
	case streamRemoteClose:
		s.state = streamClosed
		closeStream = true
		goto SEND_CLOSE

	case streamClosed:
	case streamReset:
	default:
		panic("unhandled state")
	}
	s.stateLock.Unlock()
	return nil
SEND_CLOSE:
	s.stateLock.Unlock()
	s.sendClose()
	s.notifyWaiting()
	if closeStream {
		s.session.closeStream(s.id)
	}
	return nil
}

// forceClose is used for when the session is exiting
func (s *Stream) forceClose() {
	s.stateLock.Lock()
	s.state = streamClosed
	s.stateLock.Unlock()
	s.notifyWaiting()
}

// processFlags is used to update the state of the stream
// based on set flags, if any. Lock must be held
func (s *Stream) processFlags(flags uint16) error {
	// Close the stream without holding the state lock
	closeStream := false
	defer func() {
		if closeStream {
			s.session.closeStream(s.id)
		}
	}()

	s.stateLock.Lock()
	defer s.stateLock.Unlock()
	if flags&flagACK == flagACK {
		if s.state == streamSYNSent {
			s.state = streamEstablished
		}
		s.session.establishStream(s.id)
	}
	if flags&flagFIN == flagFIN {
		switch s.state {
		case streamSYNSent:
			fallthrough
		case streamSYNReceived:
			fallthrough
		case streamEstablished:
			s.state = streamRemoteClose
			s.notifyWaiting()
		case streamLocalClose:
			s.state = streamClosed
			closeStream = true
			s.notifyWaiting()
		default:
			s.session.logger.Printf("[ERR] yamux: unexpected FIN flag in state %d", s.state)
			return ErrUnexpectedFlag
		}
	}
	if flags&flagRST == flagRST {
		s.state = streamReset
		closeStream = true
		s.notifyWaiting()
	}
	return nil
}

// notifyWaiting notifies all the waiting channels
func (s *Stream) notifyWaiting() {
	asyncNotify(s.recvNotifyCh)
	asyncNotify(s.sendNotifyCh)
}

// incrSendWindow updates the size of our send window
func (s *Stream) incrSendWindow(hdr header, flags uint16) error {
	if err := s.processFlags(flags); err != nil {
		return err
	}

	// Increase window, unblock a sender
	atomic.AddUint32(&s.sendWindow, hdr.Length())
	asyncNotify(s.sendNotifyCh)
	return nil
}

// readData is used to handle a data frame
func (s *Stream) readData(hdr header, flags uint16, conn io.Reader) error {
	if err := s.processFlags(flags); err != nil {
		return err
	}

	// Check that our recv window is not exceeded
	length := hdr.Length()
	if length == 0 {
		return nil
	}
	if remain := atomic.LoadUint32(&s.recvWindow); length > remain {
		s.session.logger.Printf("[ERR] yamux: receive window exceeded (stream: %d, remain: %d, recv: %d)", s.id, remain, length)
		return ErrRecvWindowExceeded
	}

	// Wrap in a limited reader
	conn = &io.LimitedReader{R: conn, N: int64(length)}

	// Copy into buffer
	s.recvLock.Lock()
	if s.recvBuf == nil {
		// Allocate the receive buffer just-in-time to fit the full data frame.
		// This way we can read in the whole packet without further allocations.
		s.recvBuf = bytes.NewBuffer(make([]byte, 0, length))
	}
	if _, err := io.Copy(s.recvBuf, conn); err != nil {
		s.session.logger.Printf("[ERR] yamux: Failed to read stream data: %v", err)
		s.recvLock.Unlock()
		return err
	}

	// Decrement the receive window
	atomic.AddUint32(&s.recvWindow, ^uint32(length-1))
	s.recvLock.Unlock()

	// Unblock any readers
	asyncNotify(s.recvNotifyCh)
	return nil
}

// SetDeadline sets the read and write deadlines
func (s *Stream) SetDeadline(t time.Time) error {
	if err := s.SetReadDeadline(t); err != nil {
		return err
	}
	if err := s.SetWriteDeadline(t); err != nil {
		return err
	}
	return nil
}

// SetReadDeadline sets the deadline for future Read calls.
func (s *Stream) SetReadDeadline(t time.Time) error {
	s.readDeadline = t
	return nil
}

// SetWriteDeadline sets the deadline for future Write calls
func (s *Stream) SetWriteDeadline(t time.Time) error {
	s.writeDeadline = t
	return nil
}

// Shrink is used to compact the amount of buffers utilized
// This is useful when using Yamux in a connection pool to reduce
// the idle memory utilization.
func (s *Stream) Shrink() {
	s.recvLock.Lock()
	if s.recvBuf != nil && s.recvBuf.Len() == 0 {
		s.recvBuf = nil
	}
	s.recvLock.Unlock()
}
<commit_msg>use a timer instead of 'After' to avoid leaking resources<commit_after>package yamux

import (
	"bytes"
	"io"
	"sync"
	"sync/atomic"
	"time"
)

type streamState int

const (
	streamInit streamState = iota
	streamSYNSent
	streamSYNReceived
	streamEstablished
	streamLocalClose
	streamRemoteClose
	streamClosed
	streamReset
)

// Stream is used to represent a logical stream
// within a session.
type Stream struct {
	recvWindow uint32
	sendWindow uint32

	id      uint32
	session *Session

	state     streamState
	stateLock sync.Mutex

	recvBuf  *bytes.Buffer
	recvLock sync.Mutex

	controlHdr     header
	controlErr     chan error
	controlHdrLock sync.Mutex

	sendHdr  header
	sendErr  chan error
	sendLock sync.Mutex

	recvNotifyCh chan struct{}
	sendNotifyCh chan struct{}

	readDeadline  time.Time
	writeDeadline time.Time
}

// newStream is used to construct a new stream within
// a given session for an ID
func newStream(session *Session, id uint32, state streamState) *Stream {
	s := &Stream{
		id:           id,
		session:      session,
		state:        state,
		controlHdr:   header(make([]byte, headerSize)),
		controlErr:   make(chan error, 1),
		sendHdr:      header(make([]byte, headerSize)),
		sendErr:      make(chan error, 1),
		recvWindow:   initialStreamWindow,
		sendWindow:   initialStreamWindow,
		recvNotifyCh: make(chan struct{}, 1),
		sendNotifyCh: make(chan struct{}, 1),
	}
	return s
}

// Session returns the associated stream session
func (s *Stream) Session() *Session {
	return s.session
}

// StreamID returns the ID of this stream
func (s *Stream) StreamID() uint32 {
	return s.id
}

// Read is used to read from the stream
func (s *Stream) Read(b []byte) (n int, err error) {
	defer asyncNotify(s.recvNotifyCh)
START:
	s.stateLock.Lock()
	switch s.state {
	case streamLocalClose:
		fallthrough
	case streamRemoteClose:
		fallthrough
	case streamClosed:
		if s.recvBuf == nil || s.recvBuf.Len() == 0 {
			s.stateLock.Unlock()
			return 0, io.EOF
		}
	case streamReset:
		s.stateLock.Unlock()
		return 0, ErrConnectionReset
	}
	s.stateLock.Unlock()

	// If there is no data available, block
	s.recvLock.Lock()
	if s.recvBuf == nil || s.recvBuf.Len() == 0 {
		s.recvLock.Unlock()
		goto WAIT
	}

	// Read any bytes
	n, _ = s.recvBuf.Read(b)
	s.recvLock.Unlock()

	// Send a window update potentially
	err = s.sendWindowUpdate()
	return n, err

WAIT:
	var timeout <-chan time.Time
	var timer *time.Timer
	if !s.readDeadline.IsZero() {
		delay := s.readDeadline.Sub(time.Now())
		timer = time.NewTimer(delay)
		timeout = timer.C
	}
	select {
	case <-s.recvNotifyCh:
		if timer != nil {
			timer.Stop()
		}
		goto START
	case <-timeout:
		return 0, ErrTimeout
	}
}

// Write is used to write to the stream
func (s *Stream) Write(b []byte) (n int, err error) {
	s.sendLock.Lock()
	defer s.sendLock.Unlock()
	total := 0
	for total < len(b) {
		n, err := s.write(b[total:])
		total += n
		if err != nil {
			return total, err
		}
	}
	return total, nil
}

// write is used to write to the stream, may return on
// a short write.
func (s *Stream) write(b []byte) (n int, err error) {
	var flags uint16
	var max uint32
	var body io.Reader
START:
	s.stateLock.Lock()
	switch s.state {
	case streamLocalClose:
		fallthrough
	case streamClosed:
		s.stateLock.Unlock()
		return 0, ErrStreamClosed
	case streamReset:
		s.stateLock.Unlock()
		return 0, ErrConnectionReset
	}
	s.stateLock.Unlock()

	// If there is no data available, block
	window := atomic.LoadUint32(&s.sendWindow)
	if window == 0 {
		goto WAIT
	}

	// Determine the flags if any
	flags = s.sendFlags()

	// Send up to our send window
	max = min(window, uint32(len(b)))
	body = bytes.NewReader(b[:max])

	// Send the header
	s.sendHdr.encode(typeData, flags, s.id, max)
	if err := s.session.waitForSendErr(s.sendHdr, body, s.sendErr); err != nil {
		return 0, err
	}

	// Reduce our send window
	atomic.AddUint32(&s.sendWindow, ^uint32(max-1))

	// Unlock
	return int(max), err

WAIT:
	var timeout <-chan time.Time
	if !s.writeDeadline.IsZero() {
		delay := s.writeDeadline.Sub(time.Now())
		timeout = time.After(delay)
	}
	select {
	case <-s.sendNotifyCh:
		goto START
	case <-timeout:
		return 0, ErrTimeout
	}
	return 0, nil
}

// sendFlags determines any flags that are appropriate
// based on the current stream state
func (s *Stream) sendFlags() uint16 {
	s.stateLock.Lock()
	defer s.stateLock.Unlock()
	var flags uint16
	switch s.state {
	case streamInit:
		flags |= flagSYN
		s.state = streamSYNSent
	case streamSYNReceived:
		flags |= flagACK
		s.state = streamEstablished
	}
	return flags
}

// sendWindowUpdate potentially sends a window update enabling
// further writes to take place. Must be invoked with the lock.
func (s *Stream) sendWindowUpdate() error {
	s.controlHdrLock.Lock()
	defer s.controlHdrLock.Unlock()

	// Determine the delta update
	max := s.session.config.MaxStreamWindowSize
	delta := max - atomic.LoadUint32(&s.recvWindow)

	// Determine the flags if any
	flags := s.sendFlags()

	// Check if we can omit the update
	if delta < (max/2) && flags == 0 {
		return nil
	}

	// Update our window
	atomic.AddUint32(&s.recvWindow, delta)

	// Send the header
	s.controlHdr.encode(typeWindowUpdate, flags, s.id, delta)
	if err := s.session.waitForSendErr(s.controlHdr, nil, s.controlErr); err != nil {
		return err
	}
	return nil
}

// sendClose is used to send a FIN
func (s *Stream) sendClose() error {
	s.controlHdrLock.Lock()
	defer s.controlHdrLock.Unlock()

	flags := s.sendFlags()
	flags |= flagFIN
	s.controlHdr.encode(typeWindowUpdate, flags, s.id, 0)
	if err := s.session.waitForSendErr(s.controlHdr, nil, s.controlErr); err != nil {
		return err
	}
	return nil
}

// Close is used to close the stream
func (s *Stream) Close() error {
	closeStream := false
	s.stateLock.Lock()
	switch s.state {
	// Opened means we need to signal a close
	case streamSYNSent:
		fallthrough
	case streamSYNReceived:
		fallthrough
	case streamEstablished:
		s.state = streamLocalClose
		goto SEND_CLOSE

	case streamLocalClose:
	case streamRemoteClose:
		s.state = streamClosed
		closeStream = true
		goto SEND_CLOSE

	case streamClosed:
	case streamReset:
	default:
		panic("unhandled state")
	}
	s.stateLock.Unlock()
	return nil
SEND_CLOSE:
	s.stateLock.Unlock()
	s.sendClose()
	s.notifyWaiting()
	if closeStream {
		s.session.closeStream(s.id)
	}
	return nil
}

// forceClose is used for when the session is exiting
func (s *Stream) forceClose() {
	s.stateLock.Lock()
	s.state = streamClosed
	s.stateLock.Unlock()
	s.notifyWaiting()
}

// processFlags is used to update the state of the stream
// based on set flags, if any. Lock must be held
func (s *Stream) processFlags(flags uint16) error {
	// Close the stream without holding the state lock
	closeStream := false
	defer func() {
		if closeStream {
			s.session.closeStream(s.id)
		}
	}()

	s.stateLock.Lock()
	defer s.stateLock.Unlock()
	if flags&flagACK == flagACK {
		if s.state == streamSYNSent {
			s.state = streamEstablished
		}
		s.session.establishStream(s.id)
	}
	if flags&flagFIN == flagFIN {
		switch s.state {
		case streamSYNSent:
			fallthrough
		case streamSYNReceived:
			fallthrough
		case streamEstablished:
			s.state = streamRemoteClose
			s.notifyWaiting()
		case streamLocalClose:
			s.state = streamClosed
			closeStream = true
			s.notifyWaiting()
		default:
			s.session.logger.Printf("[ERR] yamux: unexpected FIN flag in state %d", s.state)
			return ErrUnexpectedFlag
		}
	}
	if flags&flagRST == flagRST {
		s.state = streamReset
		closeStream = true
		s.notifyWaiting()
	}
	return nil
}

// notifyWaiting notifies all the waiting channels
func (s *Stream) notifyWaiting() {
	asyncNotify(s.recvNotifyCh)
	asyncNotify(s.sendNotifyCh)
}

// incrSendWindow updates the size of our send window
func (s *Stream) incrSendWindow(hdr header, flags uint16) error {
	if err := s.processFlags(flags); err != nil {
		return err
	}

	// Increase window, unblock a sender
	atomic.AddUint32(&s.sendWindow, hdr.Length())
	asyncNotify(s.sendNotifyCh)
	return nil
}

// readData is used to handle a data frame
func (s *Stream) readData(hdr header, flags uint16, conn io.Reader) error {
	if err := s.processFlags(flags); err != nil {
		return err
	}

	// Check that our recv window is not exceeded
	length := hdr.Length()
	if length == 0 {
		return nil
	}
	if remain := atomic.LoadUint32(&s.recvWindow); length > remain {
		s.session.logger.Printf("[ERR] yamux: receive window exceeded (stream: %d, remain: %d, recv: %d)", s.id, remain, length)
		return ErrRecvWindowExceeded
	}

	// Wrap in a limited reader
	conn = &io.LimitedReader{R: conn, N: int64(length)}

	// Copy into buffer
	s.recvLock.Lock()
	if s.recvBuf == nil {
		// Allocate the receive buffer just-in-time to fit the full data frame.
		// This way we can read in the whole packet without further allocations.
		s.recvBuf = bytes.NewBuffer(make([]byte, 0, length))
	}
	if _, err := io.Copy(s.recvBuf, conn); err != nil {
		s.session.logger.Printf("[ERR] yamux: Failed to read stream data: %v", err)
		s.recvLock.Unlock()
		return err
	}

	// Decrement the receive window
	atomic.AddUint32(&s.recvWindow, ^uint32(length-1))
	s.recvLock.Unlock()

	// Unblock any readers
	asyncNotify(s.recvNotifyCh)
	return nil
}

// SetDeadline sets the read and write deadlines
func (s *Stream) SetDeadline(t time.Time) error {
	if err := s.SetReadDeadline(t); err != nil {
		return err
	}
	if err := s.SetWriteDeadline(t); err != nil {
		return err
	}
	return nil
}

// SetReadDeadline sets the deadline for future Read calls.
func (s *Stream) SetReadDeadline(t time.Time) error {
	s.readDeadline = t
	return nil
}

// SetWriteDeadline sets the deadline for future Write calls
func (s *Stream) SetWriteDeadline(t time.Time) error {
	s.writeDeadline = t
	return nil
}

// Shrink is used to compact the amount of buffers utilized
// This is useful when using Yamux in a connection pool to reduce
// the idle memory utilization.
func (s *Stream) Shrink() {
	s.recvLock.Lock()
	if s.recvBuf != nil && s.recvBuf.Len() == 0 {
		s.recvBuf = nil
	}
	s.recvLock.Unlock()
}
<|endoftext|>


<commit_before>package rbd

// #cgo LDFLAGS: -lrbd
// #include <stdlib.h>
// #include <rbd/librbd.h>
import "C"

import (
	"fmt"
	"unsafe"
)

const (
	// RBD image options.

	// ImageOptionFormat is the representation of RBD_IMAGE_OPTION_FORMAT from
	// librbd
	ImageOptionFormat = C.RBD_IMAGE_OPTION_FORMAT
	// ImageOptionFeatures is the representation of RBD_IMAGE_OPTION_FEATURES
	// from librbd
	ImageOptionFeatures = C.RBD_IMAGE_OPTION_FEATURES
	// ImageOptionOrder is the representation of RBD_IMAGE_OPTION_ORDER from
	// librbd
	ImageOptionOrder = C.RBD_IMAGE_OPTION_ORDER
	// ImageOptionStripeUnit is the representation of
	// RBD_IMAGE_OPTION_STRIPE_UNIT from librbd
	ImageOptionStripeUnit = C.RBD_IMAGE_OPTION_STRIPE_UNIT
	// ImageOptionStripeCount is the representation of
	// RBD_IMAGE_OPTION_STRIPE_COUNT from librbd
	ImageOptionStripeCount = C.RBD_IMAGE_OPTION_STRIPE_COUNT
	// ImageOptionJournalOrder is the representation of
	// RBD_IMAGE_OPTION_JOURNAL_ORDER from librbd
	ImageOptionJournalOrder = C.RBD_IMAGE_OPTION_JOURNAL_ORDER
	// ImageOptionJournalSplayWidth is the representation of
	// RBD_IMAGE_OPTION_JOURNAL_SPLAY_WIDTH from librbd
	ImageOptionJournalSplayWidth = C.RBD_IMAGE_OPTION_JOURNAL_SPLAY_WIDTH
	// ImageOptionJournalPool is the representation of
	// RBD_IMAGE_OPTION_JOURNAL_POOL from librbd
	ImageOptionJournalPool = C.RBD_IMAGE_OPTION_JOURNAL_POOL
	// ImageOptionFeaturesSet is the representation of
	// RBD_IMAGE_OPTION_FEATURES_SET from librbd
	ImageOptionFeaturesSet = C.RBD_IMAGE_OPTION_FEATURES_SET
	// ImageOptionFeaturesClear is the representation of
	// RBD_IMAGE_OPTION_FEATURES_CLEAR from librbd
	ImageOptionFeaturesClear = C.RBD_IMAGE_OPTION_FEATURES_CLEAR
	// ImageOptionDataPool is the representation of RBD_IMAGE_OPTION_DATA_POOL
	// from librbd
	ImageOptionDataPool = C.RBD_IMAGE_OPTION_DATA_POOL
	// ImageOptionFlatten is the representation of RBD_IMAGE_OPTION_FLATTEN
	// from librbd
	ImageOptionFlatten = C.RBD_IMAGE_OPTION_FLATTEN
	// ImageOptionCloneFormat is the representation of
	// RBD_IMAGE_OPTION_CLONE_FORMAT from librbd
	ImageOptionCloneFormat = C.RBD_IMAGE_OPTION_CLONE_FORMAT

	// RbdImageOptionFormat deprecated alias for ImageOptionFormat
	RbdImageOptionFormat = ImageOptionFormat
	// RbdImageOptionFeatures deprecated alias for ImageOptionFeatures
	RbdImageOptionFeatures = ImageOptionFeatures
	// RbdImageOptionOrder deprecated alias for ImageOptionOrder
	RbdImageOptionOrder = ImageOptionOrder
	// RbdImageOptionStripeUnit deprecated alias for ImageOptionStripeUnit
	RbdImageOptionStripeUnit = ImageOptionStripeUnit
	// RbdImageOptionStripeCount deprecated alias for ImageOptionStripeCount
	RbdImageOptionStripeCount = ImageOptionStripeCount
	// RbdImageOptionJournalOrder deprecated alias for ImageOptionJournalOrder
	RbdImageOptionJournalOrder = ImageOptionJournalOrder
	// RbdImageOptionJournalSplayWidth deprecated alias for
	RbdImageOptionJournalSplayWidth = ImageOptionJournalSplayWidth
	// RbdImageOptionJournalPool deprecated alias for ImageOptionJournalPool
	RbdImageOptionJournalPool = ImageOptionJournalPool
	// RbdImageOptionFeaturesSet deprecated alias for ImageOptionFeaturesSet
	RbdImageOptionFeaturesSet = ImageOptionFeaturesSet
	// RbdImageOptionFeaturesClear deprecated alias for ImageOptionFeaturesClear
	RbdImageOptionFeaturesClear = ImageOptionFeaturesClear
	// RbdImageOptionDataPool deprecated alias for ImageOptionDataPool
	RbdImageOptionDataPool = ImageOptionDataPool
)

// ImageOptions represents a group of configurable image options.
type ImageOptions struct {
	options C.rbd_image_options_t
}

// ImageOption values are unique keys for configurable options.
type ImageOption C.int

// revive:disable:exported Deprecated aliases

// RbdImageOptions deprecated alias for ImageOptions
type RbdImageOptions = ImageOptions

// RbdImageOption is a deprecated alias for ImageOption
type RbdImageOption = ImageOption

//revive:enable:exported

// NewRbdImageOptions creates a new RbdImageOptions struct. Call
// RbdImageOptions.Destroy() to free the resources.
//
// Implements:
//  void rbd_image_options_create(rbd_image_options_t* opts)
func NewRbdImageOptions() *ImageOptions {
	rio := &ImageOptions{}
	C.rbd_image_options_create(&rio.options)
	return rio
}

// Destroy a RbdImageOptions struct and free the associated resources.
//
// Implements:
//  void rbd_image_options_destroy(rbd_image_options_t opts);
func (rio *ImageOptions) Destroy() {
	C.rbd_image_options_destroy(rio.options)
}

// SetString sets the value of the RbdImageOption to the given string.
//
// Implements:
//  int rbd_image_options_set_string(rbd_image_options_t opts, int optname,
//          const char* optval);
func (rio *ImageOptions) SetString(option ImageOption, value string) error {
	cValue := C.CString(value)
	defer C.free(unsafe.Pointer(cValue))

	ret := C.rbd_image_options_set_string(rio.options, C.int(option), cValue)
	if ret != 0 {
		return fmt.Errorf("%v, could not set option %v to \"%v\"",
			getError(ret), option, value)
	}

	return nil
}

// GetString returns the string value of the RbdImageOption.
//
// Implements:
//  int rbd_image_options_get_string(rbd_image_options_t opts, int optname,
//          char* optval, size_t maxlen);
func (rio *ImageOptions) GetString(option ImageOption) (string, error) {
	value := make([]byte, 4096)

	ret := C.rbd_image_options_get_string(rio.options, C.int(option),
		(*C.char)(unsafe.Pointer(&value[0])),
		C.size_t(len(value)))
	if ret != 0 {
		return "", fmt.Errorf("%v, could not get option %v", getError(ret), option)
	}

	return C.GoString((*C.char)(unsafe.Pointer(&value[0]))), nil
}

// SetUint64 sets the value of the RbdImageOption to the given uint64.
//
// Implements:
//  int rbd_image_options_set_uint64(rbd_image_options_t opts, int optname,
//          const uint64_t optval);
func (rio *ImageOptions) SetUint64(option ImageOption, value uint64) error {
	cValue := C.uint64_t(value)

	ret := C.rbd_image_options_set_uint64(rio.options, C.int(option), cValue)
	if ret != 0 {
		return fmt.Errorf("%v, could not set option %v to \"%v\"",
			getError(ret), option, value)
	}

	return nil
}

// GetUint64 returns the uint64 value of the RbdImageOption.
//
// Implements:
//  int rbd_image_options_get_uint64(rbd_image_options_t opts, int optname,
//          uint64_t* optval);
func (rio *ImageOptions) GetUint64(option ImageOption) (uint64, error) {
	var cValue C.uint64_t

	ret := C.rbd_image_options_get_uint64(rio.options, C.int(option), &cValue)
	if ret != 0 {
		return 0, fmt.Errorf("%v, could not get option %v", getError(ret), option)
	}

	return uint64(cValue), nil
}

// IsSet returns a true if the RbdImageOption is set, false otherwise.
//
// Implements:
//  int rbd_image_options_is_set(rbd_image_options_t opts, int optname,
//          bool* is_set);
func (rio *ImageOptions) IsSet(option ImageOption) (bool, error) {
	var c_set C.bool

	ret := C.rbd_image_options_is_set(rio.options, C.int(option), &c_set)
	if ret != 0 {
		return false, fmt.Errorf("%v, could not check option %v", getError(ret), option)
	}

	return bool(c_set), nil
}

// Unset a given RbdImageOption.
//
// Implements:
//  int rbd_image_options_unset(rbd_image_options_t opts, int optname)
func (rio *ImageOptions) Unset(option ImageOption) error {
	ret := C.rbd_image_options_unset(rio.options, C.int(option))
	if ret != 0 {
		return fmt.Errorf("%v, could not unset option %v", getError(ret), option)
	}

	return nil
}

// Clear all options in the RbdImageOptions.
//
// Implements:
//  void rbd_image_options_clear(rbd_image_options_t opts)
func (rio *ImageOptions) Clear() {
	C.rbd_image_options_clear(rio.options)
}

// IsEmpty returns true if there are no options set in the RbdImageOptions,
// false otherwise.
//
// Implements:
//  int rbd_image_options_is_empty(rbd_image_options_t opts)
func (rio *ImageOptions) IsEmpty() bool {
	ret := C.rbd_image_options_is_empty(rio.options)
	return ret != 0
}
<commit_msg>rbd: fix naming convention violations in options.go<commit_after>package rbd

// #cgo LDFLAGS: -lrbd
// #include <stdlib.h>
// #include <rbd/librbd.h>
import "C"

import (
	"fmt"
	"unsafe"
)

const (
	// RBD image options.

	// ImageOptionFormat is the representation of RBD_IMAGE_OPTION_FORMAT from
	// librbd
	ImageOptionFormat = C.RBD_IMAGE_OPTION_FORMAT
	// ImageOptionFeatures is the representation of RBD_IMAGE_OPTION_FEATURES
	// from librbd
	ImageOptionFeatures = C.RBD_IMAGE_OPTION_FEATURES
	// ImageOptionOrder is the representation of RBD_IMAGE_OPTION_ORDER from
	// librbd
	ImageOptionOrder = C.RBD_IMAGE_OPTION_ORDER
	// ImageOptionStripeUnit is the representation of
	// RBD_IMAGE_OPTION_STRIPE_UNIT from librbd
	ImageOptionStripeUnit = C.RBD_IMAGE_OPTION_STRIPE_UNIT
	// ImageOptionStripeCount is the representation of
	// RBD_IMAGE_OPTION_STRIPE_COUNT from librbd
	ImageOptionStripeCount = C.RBD_IMAGE_OPTION_STRIPE_COUNT
	// ImageOptionJournalOrder is the representation of
	// RBD_IMAGE_OPTION_JOURNAL_ORDER from librbd
	ImageOptionJournalOrder = C.RBD_IMAGE_OPTION_JOURNAL_ORDER
	// ImageOptionJournalSplayWidth is the representation of
	// RBD_IMAGE_OPTION_JOURNAL_SPLAY_WIDTH from librbd
	ImageOptionJournalSplayWidth = C.RBD_IMAGE_OPTION_JOURNAL_SPLAY_WIDTH
	// ImageOptionJournalPool is the representation of
	// RBD_IMAGE_OPTION_JOURNAL_POOL from librbd
	ImageOptionJournalPool = C.RBD_IMAGE_OPTION_JOURNAL_POOL
	// ImageOptionFeaturesSet is the representation of
	// RBD_IMAGE_OPTION_FEATURES_SET from librbd
	ImageOptionFeaturesSet = C.RBD_IMAGE_OPTION_FEATURES_SET
	// ImageOptionFeaturesClear is the representation of
	// RBD_IMAGE_OPTION_FEATURES_CLEAR from librbd
	ImageOptionFeaturesClear = C.RBD_IMAGE_OPTION_FEATURES_CLEAR
	// ImageOptionDataPool is the representation of RBD_IMAGE_OPTION_DATA_POOL
	// from librbd
	ImageOptionDataPool = C.RBD_IMAGE_OPTION_DATA_POOL
	// ImageOptionFlatten is the representation of RBD_IMAGE_OPTION_FLATTEN
	// from librbd
	ImageOptionFlatten = C.RBD_IMAGE_OPTION_FLATTEN
	// ImageOptionCloneFormat is the representation of
	// RBD_IMAGE_OPTION_CLONE_FORMAT from librbd
	ImageOptionCloneFormat = C.RBD_IMAGE_OPTION_CLONE_FORMAT

	// RbdImageOptionFormat deprecated alias for ImageOptionFormat
	RbdImageOptionFormat = ImageOptionFormat
	// RbdImageOptionFeatures deprecated alias for ImageOptionFeatures
	RbdImageOptionFeatures = ImageOptionFeatures
	// RbdImageOptionOrder deprecated alias for ImageOptionOrder
	RbdImageOptionOrder = ImageOptionOrder
	// RbdImageOptionStripeUnit deprecated alias for ImageOptionStripeUnit
	RbdImageOptionStripeUnit = ImageOptionStripeUnit
	// RbdImageOptionStripeCount deprecated alias for ImageOptionStripeCount
	RbdImageOptionStripeCount = ImageOptionStripeCount
	// RbdImageOptionJournalOrder deprecated alias for ImageOptionJournalOrder
	RbdImageOptionJournalOrder = ImageOptionJournalOrder
	// RbdImageOptionJournalSplayWidth deprecated alias for
	RbdImageOptionJournalSplayWidth = ImageOptionJournalSplayWidth
	// RbdImageOptionJournalPool deprecated alias for ImageOptionJournalPool
	RbdImageOptionJournalPool = ImageOptionJournalPool
	// RbdImageOptionFeaturesSet deprecated alias for ImageOptionFeaturesSet
	RbdImageOptionFeaturesSet = ImageOptionFeaturesSet
	// RbdImageOptionFeaturesClear deprecated alias for ImageOptionFeaturesClear
	RbdImageOptionFeaturesClear = ImageOptionFeaturesClear
	// RbdImageOptionDataPool deprecated alias for ImageOptionDataPool
	RbdImageOptionDataPool = ImageOptionDataPool
)

// ImageOptions represents a group of configurable image options.
type ImageOptions struct {
	options C.rbd_image_options_t
}

// ImageOption values are unique keys for configurable options.
type ImageOption C.int

// revive:disable:exported Deprecated aliases

// RbdImageOptions deprecated alias for ImageOptions
type RbdImageOptions = ImageOptions

// RbdImageOption is a deprecated alias for ImageOption
type RbdImageOption = ImageOption

//revive:enable:exported

// NewRbdImageOptions creates a new RbdImageOptions struct. Call
// RbdImageOptions.Destroy() to free the resources.
//
// Implements:
//  void rbd_image_options_create(rbd_image_options_t* opts)
func NewRbdImageOptions() *ImageOptions {
	rio := &ImageOptions{}
	C.rbd_image_options_create(&rio.options)
	return rio
}

// Destroy a RbdImageOptions struct and free the associated resources.
//
// Implements:
//  void rbd_image_options_destroy(rbd_image_options_t opts);
func (rio *ImageOptions) Destroy() {
	C.rbd_image_options_destroy(rio.options)
}

// SetString sets the value of the RbdImageOption to the given string.
//
// Implements:
//  int rbd_image_options_set_string(rbd_image_options_t opts, int optname,
//          const char* optval);
func (rio *ImageOptions) SetString(option ImageOption, value string) error {
	cValue := C.CString(value)
	defer C.free(unsafe.Pointer(cValue))

	ret := C.rbd_image_options_set_string(rio.options, C.int(option), cValue)
	if ret != 0 {
		return fmt.Errorf("%v, could not set option %v to \"%v\"",
			getError(ret), option, value)
	}

	return nil
}

// GetString returns the string value of the RbdImageOption.
//
// Implements:
//  int rbd_image_options_get_string(rbd_image_options_t opts, int optname,
//          char* optval, size_t maxlen);
func (rio *ImageOptions) GetString(option ImageOption) (string, error) {
	value := make([]byte, 4096)

	ret := C.rbd_image_options_get_string(rio.options, C.int(option),
		(*C.char)(unsafe.Pointer(&value[0])),
		C.size_t(len(value)))
	if ret != 0 {
		return "", fmt.Errorf("%v, could not get option %v", getError(ret), option)
	}

	return C.GoString((*C.char)(unsafe.Pointer(&value[0]))), nil
}

// SetUint64 sets the value of the RbdImageOption to the given uint64.
//
// Implements:
//  int rbd_image_options_set_uint64(rbd_image_options_t opts, int optname,
//          const uint64_t optval);
func (rio *ImageOptions) SetUint64(option ImageOption, value uint64) error {
	cValue := C.uint64_t(value)

	ret := C.rbd_image_options_set_uint64(rio.options, C.int(option), cValue)
	if ret != 0 {
		return fmt.Errorf("%v, could not set option %v to \"%v\"",
			getError(ret), option, value)
	}

	return nil
}

// GetUint64 returns the uint64 value of the RbdImageOption.
//
// Implements:
//  int rbd_image_options_get_uint64(rbd_image_options_t opts, int optname,
//          uint64_t* optval);
func (rio *ImageOptions) GetUint64(option ImageOption) (uint64, error) {
	var cValue C.uint64_t

	ret := C.rbd_image_options_get_uint64(rio.options, C.int(option), &cValue)
	if ret != 0 {
		return 0, fmt.Errorf("%v, could not get option %v", getError(ret), option)
	}

	return uint64(cValue), nil
}

// IsSet returns a true if the RbdImageOption is set, false otherwise.
//
// Implements:
//  int rbd_image_options_is_set(rbd_image_options_t opts, int optname,
//          bool* is_set);
func (rio *ImageOptions) IsSet(option ImageOption) (bool, error) {
	var cSet C.bool

	ret := C.rbd_image_options_is_set(rio.options, C.int(option), &cSet)
	if ret != 0 {
		return false, fmt.Errorf("%v, could not check option %v", getError(ret), option)
	}

	return bool(cSet), nil
}

// Unset a given RbdImageOption.
//
// Implements:
//  int rbd_image_options_unset(rbd_image_options_t opts, int optname)
func (rio *ImageOptions) Unset(option ImageOption) error {
	ret := C.rbd_image_options_unset(rio.options, C.int(option))
	if ret != 0 {
		return fmt.Errorf("%v, could not unset option %v", getError(ret), option)
	}

	return nil
}

// Clear all options in the RbdImageOptions.
//
// Implements:
//  void rbd_image_options_clear(rbd_image_options_t opts)
func (rio *ImageOptions) Clear() {
	C.rbd_image_options_clear(rio.options)
}

// IsEmpty returns true if there are no options set in the RbdImageOptions,
// false otherwise.
//
// Implements:
//  int rbd_image_options_is_empty(rbd_image_options_t opts)
func (rio *ImageOptions) IsEmpty() bool {
	ret := C.rbd_image_options_is_empty(rio.options)
	return ret != 0
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"

	lua "github.com/anaminus/gopher-lua"
	"github.com/anaminus/gopher-lua/parse"
	"github.com/anaminus/rbxmk"
	"github.com/anaminus/snek"
	"github.com/peterh/liner"
)

func init() {
	Program.Register(snek.Def{
		Name:    "i",
		Summary: "Enter interactive mode.",
		Description: `
Enters interactive mode. Each prompt executes a chunk of Lua code.

If a prompt begins with '=', then the comma-separated list of expressions that
follow are evaluated and printed to standard output.

The environment contains the os.exit function. When called, interactive mode is
terminated, and the program exits.

Within supported terminals, the following shortcuts are available:

	Ctrl-A, Home             Move cursor to beginning of line.
	Ctrl-E, End              Move cursor to end of line
	Ctrl-B, Left             Move cursor one character left.
	Ctrl-F, Right            Move cursor one character right.
	Ctrl-Left, Alt-B         Move cursor to previous word.
	Ctrl-Right, Alt-F        Move cursor to next word
	Ctrl-D, Del              If line is not empty, delete character under cursor.
	Ctrl-D                   If line is empty, end of file.
	Ctrl-C                   Reset input (create new empty prompt).
	Ctrl-L                   Clear screen (line is unmodified).
	Ctrl-T                   Transpose previous character with current character.
	Ctrl-H, BackSpace        Delete character before cursor.
	Ctrl-W, Alt-BackSpace    Delete word leading up to cursor.
	Alt-D                    Delete word following cursor.
	Ctrl-K                   Delete from cursor to end of line.
	Ctrl-U                   Delete from start of line to cursor.
	Ctrl-P, Up               Previous match from history.
	Ctrl-N, Down             Next match from history.
	Ctrl-R                   Reverse Search history (Ctrl-S forward, Ctrl-G cancel).
	Ctrl-Y                   Paste from Yank buffer (Alt-Y to paste next yank instead).
	Tab                      Next completion.
	Shift-Tab                (after Tab) Previous completion.
`,
		New: func() snek.Command { return &InteractiveCommand{} },
	})
}

type InteractiveCommand struct {
	WorldFlags
	Init func(rbxmk.State)
}

func (c *InteractiveCommand) SetFlags(flags snek.FlagSet) {
	c.WorldFlags.SetFlags(flags)
}

func (c *InteractiveCommand) Run(opt snek.Options) (err error) {
	// Parse flags.
	if err := opt.ParseFlags(); err != nil {
		return err
	}

	// Initialize world.
	world, err := InitWorld(WorldOpt{
		WorldFlags: c.WorldFlags,
	})
	if err != nil {
		return err
	}
	state := world.State()
	exit := make(chan struct{})
	state.GetGlobal("os").(*lua.LTable).RawSetString("exit", world.WrapFunc(func(s rbxmk.State) int {
		close(exit)
		return 0
	}))
	if c.Init != nil {
		c.Init(rbxmk.State{World: world, L: state})
	}

	// Initialize terminal prompt.
	line := liner.NewLiner()
	line.SetCtrlCAborts(true)
	line.SetMultiLineMode(true)

	// Begin read-eval-print loop.
repl:
	for {
		var chunk string
		if chunk, err = loadLine(state, line); err != nil {
			if errors.Is(err, liner.ErrPromptAborted) {
				err = nil
				break repl
			}
			if !errors.Is(err, expr) {
				break repl
			}
		}
		if chunk == "" {
			continue
		}
		if err := world.DoString(chunk, "stdin", 0); err != nil {
			fmt.Fprintln(opt.Stderr, err)
			continue
		}
		if err == expr {
			// Print values returned by chunk.
			err = nil
			n := state.GetTop()
			s := make([]interface{}, n)
			for i := 1; i <= n; i++ {
				s[i-1] = state.ToStringMeta(state.Get(i))
			}
			state.Pop(n)
			fmt.Println(s...)
		}
		// Check if os.exit was called.
		select {
		case <-exit:
			break repl
		default:
		}
	}

	if e := line.Close(); err == nil {
		return e
	}
	return err
}

// expr indicates that a chunk is an expression.
var expr = errors.New("expression")

// loadLine prompts for a Lua chunk. If the chunk begins with '=', it is
// interpreted as a return statement, and returns the expr error.
func loadLine(l *lua.LState, line *liner.State) (string, error) {
	chunk, err := line.Prompt("> ")
	if err != nil {
		return "", err
	}
	if chunk == "" {
		return "", nil
	}
	if chunk[0] == '=' {
		chunk = chunk[1:]
		if _, err := l.LoadString("return " + chunk); err == nil {
			line.AppendHistory(chunk)
			return "return " + chunk, expr
		}
	}
	if chunk, err = loadMultiline(chunk, l, line); err != nil {
		return "", err
	}
	line.AppendHistory(chunk)
	return chunk, nil
}

// loadMultiline continually prompts until a parsed Lua chunk is complete.
func loadMultiline(chunk string, l *lua.LState, line *liner.State) (string, error) {
	for {
		if _, err := l.LoadString(chunk); !isIncomplete(err) {
			return chunk, nil
		}
		line, err := line.Prompt(">> ")
		if err != nil {
			return "", err
		}
		chunk = chunk + "\n" + line
	}
}

// isIncomplete returns whether the error indicates that a parsed Lua chunk is
// incomplete.
func isIncomplete(err error) bool {
	if err == nil {
		return false
	}
	if lerr, ok := err.(*lua.ApiError); ok {
		if perr, ok := lerr.Cause.(*parse.Error); ok {
			return perr.Pos.Line == parse.EOF
		}
	}
	return false
}
<commit_msg>Fix '=' prefix missing from interactive history.<commit_after>package main

import (
	"errors"
	"fmt"

	lua "github.com/anaminus/gopher-lua"
	"github.com/anaminus/gopher-lua/parse"
	"github.com/anaminus/rbxmk"
	"github.com/anaminus/snek"
	"github.com/peterh/liner"
)

func init() {
	Program.Register(snek.Def{
		Name:    "i",
		Summary: "Enter interactive mode.",
		Description: `
Enters interactive mode. Each prompt executes a chunk of Lua code.

If a prompt begins with '=', then the comma-separated list of expressions that
follow are evaluated and printed to standard output.

The environment contains the os.exit function. When called, interactive mode is
terminated, and the program exits.

Within supported terminals, the following shortcuts are available:

	Ctrl-A, Home             Move cursor to beginning of line.
	Ctrl-E, End              Move cursor to end of line
	Ctrl-B, Left             Move cursor one character left.
	Ctrl-F, Right            Move cursor one character right.
	Ctrl-Left, Alt-B         Move cursor to previous word.
	Ctrl-Right, Alt-F        Move cursor to next word
	Ctrl-D, Del              If line is not empty, delete character under cursor.
	Ctrl-D                   If line is empty, end of file.
	Ctrl-C                   Reset input (create new empty prompt).
	Ctrl-L                   Clear screen (line is unmodified).
	Ctrl-T                   Transpose previous character with current character.
	Ctrl-H, BackSpace        Delete character before cursor.
	Ctrl-W, Alt-BackSpace    Delete word leading up to cursor.
	Alt-D                    Delete word following cursor.
	Ctrl-K                   Delete from cursor to end of line.
	Ctrl-U                   Delete from start of line to cursor.
	Ctrl-P, Up               Previous match from history.
	Ctrl-N, Down             Next match from history.
	Ctrl-R                   Reverse Search history (Ctrl-S forward, Ctrl-G cancel).
	Ctrl-Y                   Paste from Yank buffer (Alt-Y to paste next yank instead).
	Tab                      Next completion.
	Shift-Tab                (after Tab) Previous completion.
`,
		New: func() snek.Command { return &InteractiveCommand{} },
	})
}

type InteractiveCommand struct {
	WorldFlags
	Init func(rbxmk.State)
}

func (c *InteractiveCommand) SetFlags(flags snek.FlagSet) {
	c.WorldFlags.SetFlags(flags)
}

func (c *InteractiveCommand) Run(opt snek.Options) (err error) {
	// Parse flags.
	if err := opt.ParseFlags(); err != nil {
		return err
	}

	// Initialize world.
	world, err := InitWorld(WorldOpt{
		WorldFlags: c.WorldFlags,
	})
	if err != nil {
		return err
	}
	state := world.State()
	exit := make(chan struct{})
	state.GetGlobal("os").(*lua.LTable).RawSetString("exit", world.WrapFunc(func(s rbxmk.State) int {
		close(exit)
		return 0
	}))
	if c.Init != nil {
		c.Init(rbxmk.State{World: world, L: state})
	}

	// Initialize terminal prompt.
	line := liner.NewLiner()
	line.SetCtrlCAborts(true)
	line.SetMultiLineMode(true)

	// Begin read-eval-print loop.
repl:
	for {
		var chunk string
		if chunk, err = loadLine(state, line); err != nil {
			if errors.Is(err, liner.ErrPromptAborted) {
				err = nil
				break repl
			}
			if !errors.Is(err, expr) {
				break repl
			}
		}
		if chunk == "" {
			continue
		}
		if err := world.DoString(chunk, "stdin", 0); err != nil {
			fmt.Fprintln(opt.Stderr, err)
			continue
		}
		if err == expr {
			// Print values returned by chunk.
			err = nil
			n := state.GetTop()
			s := make([]interface{}, n)
			for i := 1; i <= n; i++ {
				s[i-1] = state.ToStringMeta(state.Get(i))
			}
			state.Pop(n)
			fmt.Println(s...)
		}
		// Check if os.exit was called.
		select {
		case <-exit:
			break repl
		default:
		}
	}

	if e := line.Close(); err == nil {
		return e
	}
	return err
}

// expr indicates that a chunk is an expression.
var expr = errors.New("expression")

// loadLine prompts for a Lua chunk. If the chunk begins with '=', it is
// interpreted as a return statement, and returns the expr error.
func loadLine(l *lua.LState, line *liner.State) (string, error) {
	chunk, err := line.Prompt("> ")
	if err != nil {
		return "", err
	}
	if chunk == "" {
		return "", nil
	}
	if chunk[0] == '=' {
		if _, err := l.LoadString("return " + chunk[1:]); err == nil {
			line.AppendHistory(chunk)
			return "return " + chunk[1:], expr
		}
		chunk = chunk[1:]
	}
	if chunk, err = loadMultiline(chunk, l, line); err != nil {
		return "", err
	}
	line.AppendHistory(chunk)
	return chunk, nil
}

// loadMultiline continually prompts until a parsed Lua chunk is complete.
func loadMultiline(chunk string, l *lua.LState, line *liner.State) (string, error) {
	for {
		if _, err := l.LoadString(chunk); !isIncomplete(err) {
			return chunk, nil
		}
		line, err := line.Prompt(">> ")
		if err != nil {
			return "", err
		}
		chunk = chunk + "\n" + line
	}
}

// isIncomplete returns whether the error indicates that a parsed Lua chunk is
// incomplete.
func isIncomplete(err error) bool {
	if err == nil {
		return false
	}
	if lerr, ok := err.(*lua.ApiError); ok {
		if perr, ok := lerr.Cause.(*parse.Error); ok {
			return perr.Pos.Line == parse.EOF
		}
	}
	return false
}
<|endoftext|>