<commit_before>package gowrapp

import (
	"fmt"
	"net/http"
	"reflect"
	"runtime"
	"time"

	"github.com/Sirupsen/logrus"
)

var log *logrus.Logger

type loggedResponse struct {
	w       http.ResponseWriter
	started time.Time
	status  int
	size    int
	body    string
}

func (l *loggedResponse) Flush() {
	if wf, ok := l.w.(http.Flusher); ok {
		wf.Flush()
	}
}

func (l *loggedResponse) Header() http.Header { return l.w.Header() }

func (l *loggedResponse) Write(b []byte) (int, error) {
	if l.status == 0 {
		// The status will be StatusOK if WriteHeader has not been called yet
		l.status = http.StatusOK
	}
	l.body += string(b)
	size, err := l.w.Write(b)
	l.size += size
	return size, err
}

func (l *loggedResponse) WriteHeader(status int) {
	l.w.WriteHeader(status)
	l.status = status
}

// Recover is a middleware that recovers a handler from an error and logs the traceback
func Recover(handler http.Handler) http.Handler {
	return http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			defer func() {
				if rec := recover(); rec != nil {
					var msg = "Unhandled panic: "
					var buf [4096]byte
					runtime.Stack(buf[:], true)
					stack := buf[:runtime.Stack(buf[:], false)]
					switch v := rec.(type) {
					case string:
						msg += v
					default:
						msg += reflect.TypeOf(v).String()
					}
					log.WithFields(logrus.Fields{
						"traceback": string(stack),
					}).Error(msg)
					http.Error(w, fmt.Sprintf("%s \n%s", msg, stack), http.StatusInternalServerError)
				}
			}()
			handler.ServeHTTP(w, r)
		})
}

// LogRequest is a middleware that logs a request
func LogRequest(handler http.Handler) http.Handler {
	return http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {

			lw := loggedResponse{w: w, started: time.Now()}
			handler.ServeHTTP(&lw, r)

			lm := log.WithFields(logrus.Fields{
				"status": lw.status,
				"remote": r.RemoteAddr,
				"method": r.Method,
				"proto":  r.Proto,
				"uri":    r.RequestURI,
				"took":   time.Now().Sub(lw.started),
				"size":   lw.size,
			})
			switch {
			case lw.status < 500:
				lm.Info(http.StatusText(lw.status))
			default:
				lm.Error(fmt.Sprintf("%s\n%s", http.StatusText(lw.status), lw.body))
			}
		})
}

// Error wraps http.Error so we get the error message we return logged in the system
func Error(w http.ResponseWriter, error string, code int) {
	http.Error(w, error, code)
	log.Error(error)
}

func SetLogger(mylog *logrus.Logger) {
	log = mylog
}

// RunHTTP starts a webserver with Wrapp logging and panic recovery
// The port number is fetched from the environment variable SERVICE_PORT
func RunHTTP(serviceName string, mylog *logrus.Logger, h http.Handler) {
	servicePort := GetenvDefault("SERVICE_PORT", "8080")
	SetLogger(mylog)
	log.Info(fmt.Sprintf("Starting %s on port %s", serviceName, servicePort))
	log.Fatal(http.ListenAndServe(":"+servicePort, LogRequest(Recover(h))))
}
<commit_msg>Change logging level logic, drop Error()<commit_after>package gowrapp

import (
	"fmt"
	"net/http"
	"reflect"
	"runtime"
	"time"

	"github.com/Sirupsen/logrus"
)

var log *logrus.Logger

type loggedResponse struct {
	w       http.ResponseWriter
	started time.Time
	status  int
	size    int
	body    string
}

func (l *loggedResponse) Flush() {
	if wf, ok := l.w.(http.Flusher); ok {
		wf.Flush()
	}
}

func (l *loggedResponse) Header() http.Header { return l.w.Header() }

func (l *loggedResponse) Write(b []byte) (int, error) {
	if l.status == 0 {
		// The status will be StatusOK if WriteHeader has not been called yet
		l.status = http.StatusOK
	}
	l.body += string(b)
	size, err := l.w.Write(b)
	l.size += size
	return size, err
}

func (l *loggedResponse) WriteHeader(status int) {
	l.w.WriteHeader(status)
	l.status = status
}

// Recover is a middleware that recovers a handler from an error and logs the traceback
func Recover(handler http.Handler) http.Handler {
	return http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			defer func() {
				if rec := recover(); rec != nil {
					var msg = "Unhandled panic: "
					var buf [4096]byte
					runtime.Stack(buf[:], true)
					stack := buf[:runtime.Stack(buf[:], false)]
					switch v := rec.(type) {
					case string:
						msg += v
					default:
						msg += reflect.TypeOf(v).String()
					}
					log.WithFields(logrus.Fields{
						"traceback": string(stack),
					}).Error(msg)
					http.Error(w, fmt.Sprintf("%s \n%s", msg, stack), http.StatusInternalServerError)
				}
			}()
			handler.ServeHTTP(w, r)
		})
}

// LogRequest is a middleware that logs a request
// HTTP status < 400 will be logged as Info
// HTTP status <=400 <500 will be logged as Info with the body as message
// HTTP status >= 500 will be logged as Error with the body as message
func LogRequest(handler http.Handler) http.Handler {
	return http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {

			lw := loggedResponse{w: w, started: time.Now()}
			handler.ServeHTTP(&lw, r)

			lm := log.WithFields(logrus.Fields{
				"status": lw.status,
				"remote": r.RemoteAddr,
				"method": r.Method,
				"proto":  r.Proto,
				"uri":    r.RequestURI,
				"took":   time.Now().Sub(lw.started),
				"size":   lw.size,
			})
			switch {
			case lw.status < 400:
				lm.Info(http.StatusText(lw.status))
			case lw.status >= 400 && lw.status < 500:
				lm.Info(fmt.Sprintf("%s\n%s", http.StatusText(lw.status), lw.body))
			default:
				lm.Error(fmt.Sprintf("%s\n%s", http.StatusText(lw.status), lw.body))
			}
		})
}

func SetLogger(mylog *logrus.Logger) {
	log = mylog
}

// RunHTTP starts a webserver with Wrapp logging and panic recovery
// The port number is fetched from the environment variable SERVICE_PORT
func RunHTTP(serviceName string, mylog *logrus.Logger, h http.Handler) {
	servicePort := GetenvDefault("SERVICE_PORT", "8080")
	SetLogger(mylog)
	log.Info(fmt.Sprintf("Starting %s on port %s", serviceName, servicePort))
	log.Fatal(http.ListenAndServe(":"+servicePort, LogRequest(Recover(h))))
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"mime/multipart"
	"net/http"
	"os"

	"github.com/go-martini/martini"
	"github.com/joshuarubin/goscribe/telapi"
	"github.com/kr/pretty"
	"github.com/martini-contrib/auth"
	"github.com/martini-contrib/binding"
	"github.com/martini-contrib/gzip"
	"github.com/martini-contrib/render"
	"github.com/martini-contrib/strict"
	"github.com/mitchellh/goamz/aws"
	"github.com/mitchellh/goamz/s3"
)

// TODO(jrubin):
// * mgo (http://labix.org/mgo)
// * sessions (https://github.com/martini-contrib/sessions)
// * oauth2 (https://github.com/martini-contrib/oauth2)
// * accessflags (https://github.com/martini-contrib/accessflags)
// * remove auth, replace with api key validator (github.com/martini-contrib/auth)
// * CORS support would be nice (https://github.com/martini-contrib/cors)
// * CSRF (https://github.com/martini-contrib/csrf)
// * secure (https://github.com/martini-contrib/secure)

var (
	m         *martini.ClassicMartini
	s3Bucket  *s3.Bucket
	s3BaseURL string
	baseURL   string
	authUser  string
	authPass  string
)

type callbackData struct {
	CallbackURL string `form:"callback_url" binding:"required"`
}

type transcribeData struct {
	callbackData
	AudioURL string `form:"audio_url" binding:"required"`
}

type transcribeUploadData struct {
	callbackData
	AudioData *multipart.FileHeader `form:"audio_data" binding:"required"`
}

func init() {
	// BASE_URL, AUTH_USER and AUTH_PASS, AWS_S3_BASE_URL are not required or else wercker tests would fail
	baseURL = os.Getenv("BASE_URL")
	authUser = os.Getenv("AUTH_USER")
	authPass = os.Getenv("AUTH_PASS")
	s3BaseURL = os.Getenv("AWS_S3_BASE_URL")

	if awsAuth, err := aws.EnvAuth(); err != nil {
		// not required or else wercker tests would fail
		log.Println(err)
	} else {
		// TODO(jrubin) allow region to be chosen by env variable
		s3Data := s3.New(awsAuth, aws.USWest2)
		s3Bucket = s3Data.Bucket(os.Getenv("AWS_S3_BUCKET_NAME"))
	}

	m = martini.Classic()

	m.Use(gzip.All())
	m.Use(render.Renderer())

	m.Get("/", func() string {
		return "hello, world"
	})

	m.Post(
		"/v1/transcribe",
		auth.Basic(authUser, authPass),
		strict.Accept("application/json"),
		strict.ContentType("application/x-www-form-urlencoded"),
		binding.Bind(transcribeData{}),
		binding.ErrorHandler,
		handleTranscribe,
	)

	m.Post(
		"/v1/transcribe/process",
		strict.ContentType("application/x-www-form-urlencoded"),
		binding.Bind(telapi.TranscribeCallbackData{}),
		binding.ErrorHandler,
		handleTranscribeProcess,
	)

	m.Post(
		"/v1/transcribe/upload",
		auth.Basic(authUser, authPass),
		strict.Accept("application/json"),
		binding.MultipartForm(transcribeUploadData{}),
		binding.ErrorHandler,
		handleTranscribeUpload,
	)

	m.Router.NotFound(strict.MethodNotAllowed, strict.NotFound)
}

func main() {
	m.Run()
}

func success(r render.Render, data ...interface{}) {
	ret := map[string]interface{}{
		"status": http.StatusOK,
	}

	l := len(data)
	if l == 1 {
		ret["data"] = data[0]
	} else if l > 1 {
		ret["data"] = data
	}

	r.JSON(http.StatusOK, ret)
}

func jsonError(r render.Render, status int, err error) {
	r.JSON(status, map[string]interface{}{
		"status": status,
		"error":  err.Error(),
	})
}

func telapiError(r render.Render, err error) {
	if telapiError, ok := err.(telapi.Error); ok {
		r.JSON(telapiError.JSON())
		return
	}

	jsonError(r, http.StatusInternalServerError, err)
}

func s3Upload(basePath string, data *multipart.FileHeader) (string, error) {
	fileName := data.Filename
	fileType := data.Header.Get("Content-Type")

	file, err := data.Open()
	defer file.Close()

	if err != nil {
		return "", err
	}

	fileData, err := ioutil.ReadAll(file)
	if err != nil {
		return "", err
	}

	filePath := fmt.Sprintf("/%s/%s", basePath, fileName)

	if err := s3Bucket.Put(filePath, fileData, fileType, s3.PublicRead); err != nil {
		return "", err
	}

	return s3BaseURL + filePath, nil
}

func handleTranscribe(data transcribeData, r render.Render) {
	resp, err := telapi.TranscribeURL(data.AudioURL, data.CallbackURL)
	if err != nil {
		telapiError(r, err)
		return
	}

	success(r, resp.TranscribeClientResponse.Translate())
}

func handleTranscribeProcess(data telapi.TranscribeCallbackData, r render.Render) {
	pretty.Println(data)
	b, _ := json.Marshal(data.TranscribeCallbackClientData)
	pretty.Println(string(b))
	success(r)
}

func handleTranscribeUpload(data transcribeUploadData, r render.Render) {
	fileURL, err := s3Upload("audio", data.AudioData)
	if err != nil {
		jsonError(r, http.StatusInternalServerError, err)
		return
	}

	resp, err := telapi.TranscribeURL(fileURL, data.CallbackURL)
	if err != nil {
		telapiError(r, err)
		return
	}

	success(r, resp.TranscribeClientResponse.Translate())
}
<commit_msg>cleanup web server<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"mime/multipart"
	"net/http"
	"os"

	"github.com/go-martini/martini"
	"github.com/joshuarubin/goscribe/telapi"
	"github.com/kr/pretty"
	"github.com/martini-contrib/auth"
	"github.com/martini-contrib/binding"
	"github.com/martini-contrib/gzip"
	"github.com/martini-contrib/render"
	"github.com/martini-contrib/strict"
	"github.com/mitchellh/goamz/aws"
	"github.com/mitchellh/goamz/s3"
)

// TODO(jrubin):
// * mgo (http://labix.org/mgo)
// * sessions (https://github.com/martini-contrib/sessions)
// * oauth2 (https://github.com/martini-contrib/oauth2)
// * accessflags (https://github.com/martini-contrib/accessflags)
// * remove auth, replace with api key validator (github.com/martini-contrib/auth)
// * CORS support would be nice (https://github.com/martini-contrib/cors)
// * CSRF (https://github.com/martini-contrib/csrf)
// * secure (https://github.com/martini-contrib/secure)

var (
	m         *martini.ClassicMartini
	s3Bucket  *s3.Bucket
	s3BaseURL string
	baseURL   string
	authUser  string
	authPass  string
)

type callbackData struct {
	CallbackURL string `form:"callback_url" binding:"required"`
}

type transcribeData struct {
	callbackData
	AudioURL string `form:"audio_url" binding:"required"`
}

type transcribeUploadData struct {
	callbackData
	AudioData *multipart.FileHeader `form:"audio_data" binding:"required"`
}

func init() {
	// BASE_URL, AUTH_USER and AUTH_PASS, AWS_S3_BASE_URL are not required or else wercker tests would fail
	baseURL = os.Getenv("BASE_URL")
	authUser = os.Getenv("AUTH_USER")
	authPass = os.Getenv("AUTH_PASS")
	s3BaseURL = os.Getenv("AWS_S3_BASE_URL")

	if awsAuth, err := aws.EnvAuth(); err != nil {
		// not required or else wercker tests would fail
		log.Println(err)
	} else {
		// TODO(jrubin) allow region to be chosen by env variable
		s3Data := s3.New(awsAuth, aws.USWest2)
		s3Bucket = s3Data.Bucket(os.Getenv("AWS_S3_BUCKET_NAME"))
	}

	m = martini.Classic()

	m.Use(gzip.All())
	m.Use(render.Renderer())

	m.Get("/", func() string {
		return "hello, world"
	})

	m.Post(
		"/v1/transcribe",
		auth.Basic(authUser, authPass),
		strict.Accept("application/json"),
		strict.ContentType("application/x-www-form-urlencoded"),
		binding.Bind(transcribeData{}),
		binding.ErrorHandler,
		handleTranscribe,
	)

	m.Post(
		"/v1/transcribe/process",
		strict.ContentType("application/x-www-form-urlencoded"),
		binding.Bind(telapi.TranscribeCallbackData{}),
		binding.ErrorHandler,
		handleTranscribeProcess,
	)

	m.Post(
		"/v1/transcribe/upload",
		auth.Basic(authUser, authPass),
		strict.Accept("application/json"),
		binding.MultipartForm(transcribeUploadData{}),
		binding.ErrorHandler,
		handleTranscribeUpload,
	)

	m.Router.NotFound(strict.MethodNotAllowed, strict.NotFound)
}

func main() {
	m.Run()
}

func success(r render.Render, data ...interface{}) {
	ret := map[string]interface{}{
		"status": http.StatusOK,
	}

	l := len(data)
	if l == 1 {
		ret["data"] = data[0]
	} else if l > 1 {
		ret["data"] = data
	}

	r.JSON(http.StatusOK, ret)
}

func jsonError(r render.Render, status int, err error) {
	r.JSON(status, map[string]interface{}{
		"status": status,
		"error":  err.Error(),
	})
}

func telapiError(r render.Render, err error) {
	if telapiError, ok := err.(telapi.Error); ok {
		r.JSON(telapiError.JSON())
		return
	}

	jsonError(r, http.StatusInternalServerError, err)
}

func transcribe(r render.Render, url, callback string) {
	resp, err := telapi.TranscribeURL(url, callback)
	if err != nil {
		telapiError(r, err)
		return
	}

	success(r, resp.TranscribeClientResponse.Translate())
}

func s3Upload(basePath string, data *multipart.FileHeader) (string, error) {
	fileName := data.Filename
	fileType := data.Header.Get("Content-Type")

	file, err := data.Open()
	defer file.Close()

	if err != nil {
		return "", err
	}

	fileData, err := ioutil.ReadAll(file)
	if err != nil {
		return "", err
	}

	filePath := fmt.Sprintf("/%s/%s", basePath, fileName)

	if err := s3Bucket.Put(filePath, fileData, fileType, s3.PublicRead); err != nil {
		return "", err
	}

	return s3BaseURL + filePath, nil
}

func handleTranscribe(data transcribeData, r render.Render) {
	transcribe(r, data.AudioURL, data.CallbackURL)
}

func handleTranscribeProcess(data telapi.TranscribeCallbackData, r render.Render) {
	pretty.Println(data)
	b, _ := json.Marshal(data.TranscribeCallbackClientData)
	pretty.Println(string(b))
	success(r)
}

func handleTranscribeUpload(data transcribeUploadData, r render.Render) {
	fileURL, err := s3Upload("audio", data.AudioData)
	if err != nil {
		jsonError(r, http.StatusInternalServerError, err)
		return
	}

	transcribe(r, fileURL, data.CallbackURL)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/gin-gonic/gin"
)

func InitWeb(logLevel, assets string) *gin.Engine {
	if logLevel != "debug" {
		gin.SetMode(gin.ReleaseMode)
	}

	engine := gin.New()
	engine.HandleMethodNotAllowed = true
	engine.Use(gin.Recovery())
	engine.Static("/static", assets)

	engine.Use(func(c *gin.Context) {
		start := time.Now()

		c.Next()

		end := time.Now()
		latency := end.Sub(start)
		clientIP := c.Request.RemoteAddr
		method := c.Request.Method
		statusCode := c.Writer.Status()
		errors := ""
		if c.Errors != nil {
			errors = fmt.Sprintf(" - %s", c.Errors.String())
		}

		log.Infof("[GIN] %v - %3d - %12v - %s %s %s%s",
			end.Format("2006-01-02 15:04:05"),
			statusCode,
			latency,
			clientIP,
			method,
			c.Request.URL.Path,
			errors,
		)
	})

	engine.Use(func(c *gin.Context) {
		c.Writer.Header().Set("Server", "SMS Logger")
		c.Next()
	})

	return engine
}

func jsonWebError(c *gin.Context, status int, message string, args ...interface{}) {
	log.Errorf("Error (%d): %s", status, fmt.Sprintf(message, args...))
	c.JSON(status, gin.H{"output": fmt.Sprintf(message, args...)})
	c.Abort()
}

func stringWebError(c *gin.Context, status int, message string, args ...interface{}) {
	log.Errorf("Error (%d): %s", status, fmt.Sprintf(message, args...))
	c.String(status, message, args)
	c.Abort()
}
<commit_msg>Redirect trailing slash<commit_after>package main

import (
	"fmt"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/gin-gonic/gin"
)

func InitWeb(logLevel, assets string) *gin.Engine {
	if logLevel != "debug" {
		gin.SetMode(gin.ReleaseMode)
	}

	engine := gin.New()
	engine.HandleMethodNotAllowed = true
	engine.RedirectTrailingSlash = true
	engine.Use(gin.Recovery())
	engine.Static("/static", assets)

	engine.Use(func(c *gin.Context) {
		start := time.Now()

		c.Next()

		end := time.Now()
		latency := end.Sub(start)
		clientIP := c.Request.RemoteAddr
		method := c.Request.Method
		statusCode := c.Writer.Status()
		errors := ""
		if c.Errors != nil {
			errors = fmt.Sprintf(" - %s", c.Errors.String())
		}

		log.Infof("[GIN] %v - %3d - %12v - %s %s %s%s",
			end.Format("2006-01-02 15:04:05"),
			statusCode,
			latency,
			clientIP,
			method,
			c.Request.URL.Path,
			errors,
		)
	})

	engine.Use(func(c *gin.Context) {
		c.Writer.Header().Set("Server", "SMS Logger")
		c.Next()
	})

	return engine
}

func jsonWebError(c *gin.Context, status int, message string, args ...interface{}) {
	log.Errorf("Error (%d): %s", status, fmt.Sprintf(message, args...))
	c.JSON(status, gin.H{"output": fmt.Sprintf(message, args...)})
	c.Abort()
}

func stringWebError(c *gin.Context, status int, message string, args ...interface{}) {
	log.Errorf("Error (%d): %s", status, fmt.Sprintf(message, args...))
	c.String(status, message, args)
	c.Abort()
}
<|endoftext|>


<commit_before>package pb_test

import (
	"fmt"

	"github.com/bigwhite/gowechat/qy"
)

const (
	accessToken = "wx1234abcd"
	layout      = `{
	    "button": [
			{
				"name":"submenu1",
				"subbutton": [
				    {
						"name": "item1",
						"type":"click",
						"key":"s1-item1"
					},
				    {
						"name": "item2",
						"type":"click",
						"key":"s1-item2"
					}
				]
			},
			{
				"name":"submenu2",
				"subbutton": [
				    {
						"name": "item1",
						"type":"click",
						"key":"s2-item1"
					},
				    {
						"name": "item2",
						"type":"click",
						"key":"s2-item2"
					}
				]
			}
		]
} `
	agentID = "5"
)

func ExampleCreateMenu() {
	err := qy.CreateMenu([]byte(layout), accessToken, agentID)
	fmt.Println(err)
	// Output: https://qyapi.weixin.qq.com/cgi-bin/menu/create?access_token=wx1234abcd&agentid=5
	// invalid access_token
}
<commit_msg>move example_test.go to qy<commit_after><|endoftext|>


<commit_before>package agent

import (
	"context"
	"fmt"
	"time"

	"github.com/coreos/go-systemd/login1"
	"github.com/golang/glog"
	v1meta "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/kubernetes"
	v1core "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/pkg/api/v1"

	"github.com/coreos/container-linux-update-operator/pkg/constants"
	"github.com/coreos/container-linux-update-operator/pkg/drain"
	"github.com/coreos/container-linux-update-operator/pkg/k8sutil"
	"github.com/coreos/container-linux-update-operator/pkg/updateengine"
)

type Klocksmith struct {
	node string
	kc   *kubernetes.Clientset
	nc   v1core.NodeInterface
	ue   *updateengine.Client
	lc   *login1.Conn
}

var (
	shouldRebootSelector = fields.Set(map[string]string{
		constants.AnnotationOkToReboot:   constants.True,
		constants.AnnotationRebootNeeded: constants.True,
	}).AsSelector()
)

func New(node string) (*Klocksmith, error) {
	// set up kubernetes in-cluster client
	kc, err := k8sutil.InClusterClient()
	if err != nil {
		return nil, fmt.Errorf("error creating kubernetes client: %v", err)
	}

	// node interface
	nc := kc.Nodes()

	// set up update_engine client
	ue, err := updateengine.New()
	if err != nil {
		return nil, fmt.Errorf("error establishing connection to update_engine dbus: %v", err)
	}

	// set up login1 client for our eventual reboot
	lc, err := login1.New()
	if err != nil {
		return nil, fmt.Errorf("error establishing connection to logind dbus: %v", err)
	}

	return &Klocksmith{node, kc, nc, ue, lc}, nil
}

// updateStatusCallback receives Status messages from update engine. If the
// status is UpdateStatusUpdatedNeedReboot, indicate that with a label on our
// node.
func (k *Klocksmith) updateStatusCallback(s updateengine.Status) {
	glog.Info("Updating status")
	// update our status
	anno := map[string]string{
		constants.AnnotationStatus:          s.CurrentOperation,
		constants.AnnotationLastCheckedTime: fmt.Sprintf("%d", s.LastCheckedTime),
		constants.AnnotationNewVersion:      s.NewVersion,
	}

	// indicate we need a reboot
	if s.CurrentOperation == updateengine.UpdateStatusUpdatedNeedReboot {
		glog.Info("Indicating a reboot is needed")
		anno[constants.AnnotationRebootNeeded] = constants.True
	}

	wait.PollUntil(10*time.Second, func() (bool, error) {
		if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
			glog.Errorf("Failed to set annotation %q: %v", constants.AnnotationStatus, err)
			return false, nil
		}

		return true, nil
	}, wait.NeverStop)
}

// setInfoLabels labels our node with helpful info about Container Linux.
func (k *Klocksmith) setInfoLabels() error {
	vi, err := k8sutil.GetVersionInfo()
	if err != nil {
		return fmt.Errorf("failed to get version info: %v", err)
	}

	labels := map[string]string{
		constants.LabelID:      vi.ID,
		constants.LabelGroup:   vi.Group,
		constants.LabelVersion: vi.Version,
	}

	if err := k8sutil.SetNodeLabels(k.nc, k.node, labels); err != nil {
		return err
	}

	return nil
}

// Run runs klocksmithd, reacting to the update_engine reboot signal by
// draining pods on this kubernetes node and rebooting.
//
// TODO(mischief): try to be more resilient against transient failures
func (k *Klocksmith) Run() error {
	glog.Info("Setting info labels")
	if err := k.setInfoLabels(); err != nil {
		return fmt.Errorf("failed to set node info: %v", err)
	}

	// we are schedulable now.
	glog.Info("Marking node as schedulable")
	if err := k8sutil.Unschedulable(k.nc, k.node, false); err != nil {
		return err
	}

	// set coreos.com/update1/reboot-in-progress=false and
	// coreos.com/update1/reboot-needed=false
	anno := map[string]string{
		constants.AnnotationRebootInProgress: constants.False,
		constants.AnnotationRebootNeeded:     constants.False,
	}
	glog.Infof("Setting annotations %#v", anno)
	if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
		return err
	}
	// Since we set 'reboot-needed=false', 'ok-to-reboot' should clear.
	// Wait for it to do so, else we might start reboot-looping
	if err := k.waitForNotOkToReboot(); err != nil {
		return err
	}

	// watch update engine for status updates
	go k.watchUpdateStatus(k.updateStatusCallback)

	// block until constants.AnnotationOkToReboot is set
	glog.Infof("Waiting for ok-to-reboot from controller...")
	if err := k.waitForOkToReboot(); err != nil {
		return err
	}

	// set constants.AnnotationRebootInProgress and drain self
	anno = map[string]string{
		constants.AnnotationRebootInProgress: constants.True,
	}

	glog.Infof("Setting annotations %#v", anno)
	if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
		return err
	}

	// drain self equates to:
	// 1. set Unschedulable
	// 2. delete all pods
	// unlike `kubectl drain`, we do not care about emptyDir or orphan pods
	// ('any pods that are neither mirror pods nor managed by
	// ReplicationController, ReplicaSet, DaemonSet or Job')

	glog.Info("Marking node as unschedulable")
	if err := k8sutil.Unschedulable(k.nc, k.node, true); err != nil {
		return err
	}

	glog.Info("Getting pod list for deletion")
	pods, err := k.getPodsForDeletion()
	if err != nil {
		return err
	}

	// delete the pods.
	// TODO(mischief): explicitly don't terminate self? we'll probably just be a
	// mirror pod or daemonset anyway..
	glog.Infof("Deleting %d pods", len(pods))
	deleteOptions := v1meta.NewDeleteOptions(30)
	for _, pod := range pods {
		glog.Infof("Terminating pod %q...", pod.Name)
		if err := k.kc.Pods(pod.Namespace).Delete(pod.Name, deleteOptions); err != nil {
			return fmt.Errorf("failed terminating pod %q: %v", pod.Name, err)
		}
	}

	glog.Info("Node drained, rebooting")

	// reboot
	k.lc.Reboot(false)

	// cross fingers
	time.Sleep(24 * 7 * time.Hour)

	return nil
}

func (k *Klocksmith) watchUpdateStatus(update func(s updateengine.Status)) {
	glog.Info("Beginning to watch update_engine status")

	oldOperation := ""
	ch := make(chan updateengine.Status, 1)
	stop := make(chan struct{})

	go k.ue.ReceiveStatuses(ch, stop)

	// synthesize first message
	st, err := k.ue.GetStatus()
	if err != nil {
		glog.Fatalf("Failed to get update_engine status: %v", err)
	}

	ch <- st

	for status := range ch {
		if status.CurrentOperation != oldOperation && update != nil {
			update(status)
			oldOperation = status.CurrentOperation
		}
	}
}

// waitForRebootSignal watches update_engine and waits for
// updateengine.UpdateStatusUpdatedNeedReboot
func (k *Klocksmith) waitForRebootSignal() error {
	// first, check current status.
	status, err := k.ue.GetStatus()
	if err != nil {
		return fmt.Errorf("failed to get update_engine status: %v", err)
	}

	// if we're not in UpdateStatusUpdatedNeedReboot, then wait for it.
	if status.CurrentOperation != updateengine.UpdateStatusUpdatedNeedReboot {
		ctx, cancel := context.WithCancel(context.Background())
		stc := make(chan updateengine.Status)
		go k.ue.RebootNeededSignal(stc, ctx.Done())
		<-stc
		cancel()
	}

	// we are now in UpdateStatusUpdatedNeedReboot state.
	return nil
}

// waitForOkToReboot waits for both 'ok-to-reboot' and 'needs-reboot' to be true.
func (k *Klocksmith) waitForOkToReboot() error {
	n, err := k.nc.Get(k.node, v1meta.GetOptions{})
	if err != nil {
		return fmt.Errorf("failed to get self node (%q): %v", k.node, err)
	}

	if n.Annotations[constants.AnnotationOkToReboot] == constants.True && n.Annotations[constants.AnnotationRebootNeeded] == constants.True {
		return nil
	}

	// XXX: set timeout > 0?
	watcher, err := k.nc.Watch(v1meta.ListOptions{
		FieldSelector:   fields.OneTermEqualSelector("metadata.name", n.Name).String(),
		ResourceVersion: n.ResourceVersion,
	})
	if err != nil {
		return fmt.Errorf("failed to watch self node (%q): %v", k.node, err)
	}

	// hopefully 24 hours is enough time between indicating we need a
	// reboot and the controller telling us to do it
	ev, err := watch.Until(time.Hour*24, watcher, k8sutil.NodeAnnotationCondition(shouldRebootSelector))
	if err != nil {
		return fmt.Errorf("waiting for annotation %q failed: %v", constants.AnnotationOkToReboot, err)
	}

	// sanity check
	no, ok := ev.Object.(*v1.Node)
	if !ok {
		panic("event contains a non-*api.Node object")
	}

	if no.Annotations[constants.AnnotationOkToReboot] != constants.True {
		panic("event did not contain annotation expected")
	}

	return nil
}

func (k *Klocksmith) waitForNotOkToReboot() error {
	n, err := k.nc.Get(k.node, v1meta.GetOptions{})
	if err != nil {
		return fmt.Errorf("failed to get self node (%q): %v", k.node, err)
	}

	if n.Annotations[constants.AnnotationOkToReboot] != constants.True {
		return nil
	}

	// XXX: set timeout > 0?
	watcher, err := k.nc.Watch(v1meta.ListOptions{
		FieldSelector:   fields.OneTermEqualSelector("metadata.name", n.Name).String(),
		ResourceVersion: n.ResourceVersion,
	})
	if err != nil {
		return fmt.Errorf("failed to watch self node (%q): %v", k.node, err)
	}

	// Within 24 hours of indicating we don't need a reboot we should be given a not-ok.
	// If that isn't the case, it likely means the operator isn't running, and
	// we'll just crash-loop in that case, and hopefully that will help the user realize something's wrong.
	// Use a custom condition function to use the more correct 'OkToReboot !=
	// true' vs '== False'; due to the operator matching on '== True', and not
	// going out of its way to convert '' => 'False', checking the exact inverse
	// of what the operator checks is the correct thing to do.
	ev, err := watch.Until(time.Hour*24, watcher, watch.ConditionFunc(func(event watch.Event) (bool, error) {
		switch event.Type {
		case watch.Error:
			return false, fmt.Errorf("error watching node: %v", event.Object)
		case watch.Deleted:
			return false, fmt.Errorf("our node was deleted while we were waiting for ready")
		}

		no := event.Object.(*v1.Node)
		if no.Annotations[constants.AnnotationOkToReboot] != constants.True {
			return true, nil
		}
		return false, nil
	}))
	if err != nil {
		return fmt.Errorf("waiting for annotation %q failed: %v", constants.AnnotationOkToReboot, err)
	}

	// sanity check
	no := ev.Object.(*v1.Node)

	if no.Annotations[constants.AnnotationOkToReboot] == constants.True {
		panic("event did not contain annotation expected")
	}

	return nil
}

func (k *Klocksmith) getPodsForDeletion() ([]v1.Pod, error) {
	pods, err := drain.GetPodsForDeletion(k.kc, k.node)
	if err != nil {
		return nil, fmt.Errorf("failed to get list of pods for deletion: %v", err)
	}

	// XXX: ignoring kube-system is a simple way to avoid eviciting
	// critical components such as kube-scheduler and
	// kube-controller-manager.

	pods = k8sutil.FilterPods(pods, func(p *v1.Pod) bool {
		if p.Namespace == "kube-system" {
			return false
		}

		return true
	})

	return pods, nil
}
<commit_msg>pkg/agent: avoid exiting on watch termination<commit_after>package agent

import (
	"context"
	"fmt"
	"time"

	"github.com/coreos/go-systemd/login1"
	"github.com/golang/glog"
	v1meta "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/kubernetes"
	v1core "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/pkg/api/v1"

	"github.com/coreos/container-linux-update-operator/pkg/constants"
	"github.com/coreos/container-linux-update-operator/pkg/drain"
	"github.com/coreos/container-linux-update-operator/pkg/k8sutil"
	"github.com/coreos/container-linux-update-operator/pkg/updateengine"
)

type Klocksmith struct {
	node string
	kc   *kubernetes.Clientset
	nc   v1core.NodeInterface
	ue   *updateengine.Client
	lc   *login1.Conn
}

var (
	shouldRebootSelector = fields.Set(map[string]string{
		constants.AnnotationOkToReboot:   constants.True,
		constants.AnnotationRebootNeeded: constants.True,
	}).AsSelector()
)

func New(node string) (*Klocksmith, error) {
	// set up kubernetes in-cluster client
	kc, err := k8sutil.InClusterClient()
	if err != nil {
		return nil, fmt.Errorf("error creating kubernetes client: %v", err)
	}

	// node interface
	nc := kc.Nodes()

	// set up update_engine client
	ue, err := updateengine.New()
	if err != nil {
		return nil, fmt.Errorf("error establishing connection to update_engine dbus: %v", err)
	}

	// set up login1 client for our eventual reboot
	lc, err := login1.New()
	if err != nil {
		return nil, fmt.Errorf("error establishing connection to logind dbus: %v", err)
	}

	return &Klocksmith{node, kc, nc, ue, lc}, nil
}

// updateStatusCallback receives Status messages from update engine. If the
// status is UpdateStatusUpdatedNeedReboot, indicate that with a label on our
// node.
func (k *Klocksmith) updateStatusCallback(s updateengine.Status) {
	glog.Info("Updating status")
	// update our status
	anno := map[string]string{
		constants.AnnotationStatus:          s.CurrentOperation,
		constants.AnnotationLastCheckedTime: fmt.Sprintf("%d", s.LastCheckedTime),
		constants.AnnotationNewVersion:      s.NewVersion,
	}

	// indicate we need a reboot
	if s.CurrentOperation == updateengine.UpdateStatusUpdatedNeedReboot {
		glog.Info("Indicating a reboot is needed")
		anno[constants.AnnotationRebootNeeded] = constants.True
	}

	wait.PollUntil(10*time.Second, func() (bool, error) {
		if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
			glog.Errorf("Failed to set annotation %q: %v", constants.AnnotationStatus, err)
			return false, nil
		}

		return true, nil
	}, wait.NeverStop)
}

// setInfoLabels labels our node with helpful info about Container Linux.
func (k *Klocksmith) setInfoLabels() error {
	vi, err := k8sutil.GetVersionInfo()
	if err != nil {
		return fmt.Errorf("failed to get version info: %v", err)
	}

	labels := map[string]string{
		constants.LabelID:      vi.ID,
		constants.LabelGroup:   vi.Group,
		constants.LabelVersion: vi.Version,
	}

	if err := k8sutil.SetNodeLabels(k.nc, k.node, labels); err != nil {
		return err
	}

	return nil
}

// Run runs klocksmithd, reacting to the update_engine reboot signal by
// draining pods on this kubernetes node and rebooting.
//
// TODO(mischief): try to be more resilient against transient failures
func (k *Klocksmith) Run() error {
	glog.Info("Setting info labels")
	if err := k.setInfoLabels(); err != nil {
		return fmt.Errorf("failed to set node info: %v", err)
	}

	// we are schedulable now.
	glog.Info("Marking node as schedulable")
	if err := k8sutil.Unschedulable(k.nc, k.node, false); err != nil {
		return err
	}

	// set coreos.com/update1/reboot-in-progress=false and
	// coreos.com/update1/reboot-needed=false
	anno := map[string]string{
		constants.AnnotationRebootInProgress: constants.False,
		constants.AnnotationRebootNeeded:     constants.False,
	}
	glog.Infof("Setting annotations %#v", anno)
	if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
		return err
	}
	// Since we set 'reboot-needed=false', 'ok-to-reboot' should clear.
	// Wait for it to do so, else we might start reboot-looping
	if err := k.waitForNotOkToReboot(); err != nil {
		return err
	}

	// watch update engine for status updates
	go k.watchUpdateStatus(k.updateStatusCallback)

	// block until constants.AnnotationOkToReboot is set
	for {
		glog.Infof("Waiting for ok-to-reboot from controller...")
		err := k.waitForOkToReboot()
		if err == nil {
			// time to reboot
			break
		}
		glog.Warningf("error waiting for an ok-to-reboot: %v", err)
	}

	// set constants.AnnotationRebootInProgress and drain self
	anno = map[string]string{
		constants.AnnotationRebootInProgress: constants.True,
	}

	glog.Infof("Setting annotations %#v", anno)
	if err := k8sutil.SetNodeAnnotations(k.nc, k.node, anno); err != nil {
		return err
	}

	// drain self equates to:
	// 1. set Unschedulable
	// 2. delete all pods
	// unlike `kubectl drain`, we do not care about emptyDir or orphan pods
	// ('any pods that are neither mirror pods nor managed by
	// ReplicationController, ReplicaSet, DaemonSet or Job')

	glog.Info("Marking node as unschedulable")
	if err := k8sutil.Unschedulable(k.nc, k.node, true); err != nil {
		return err
	}

	glog.Info("Getting pod list for deletion")
	pods, err := k.getPodsForDeletion()
	if err != nil {
		return err
	}

	// delete the pods.
	// TODO(mischief): explicitly don't terminate self? we'll probably just be a
	// mirror pod or daemonset anyway..
	glog.Infof("Deleting %d pods", len(pods))
	deleteOptions := v1meta.NewDeleteOptions(30)
	for _, pod := range pods {
		glog.Infof("Terminating pod %q...", pod.Name)
		if err := k.kc.Pods(pod.Namespace).Delete(pod.Name, deleteOptions); err != nil {
			return fmt.Errorf("failed terminating pod %q: %v", pod.Name, err)
		}
	}

	glog.Info("Node drained, rebooting")

	// reboot
	k.lc.Reboot(false)

	// cross fingers
	time.Sleep(24 * 7 * time.Hour)

	return nil
}

func (k *Klocksmith) watchUpdateStatus(update func(s updateengine.Status)) {
	glog.Info("Beginning to watch update_engine status")

	oldOperation := ""
	ch := make(chan updateengine.Status, 1)
	stop := make(chan struct{})

	go k.ue.ReceiveStatuses(ch, stop)

	// synthesize first message
	st, err := k.ue.GetStatus()
	if err != nil {
		glog.Fatalf("Failed to get update_engine status: %v", err)
	}

	ch <- st

	for status := range ch {
		if status.CurrentOperation != oldOperation && update != nil {
			update(status)
			oldOperation = status.CurrentOperation
		}
	}
}

// waitForRebootSignal watches update_engine and waits for
// updateengine.UpdateStatusUpdatedNeedReboot
func (k *Klocksmith) waitForRebootSignal() error {
	// first, check current status.
	status, err := k.ue.GetStatus()
	if err != nil {
		return fmt.Errorf("failed to get update_engine status: %v", err)
	}

	// if we're not in UpdateStatusUpdatedNeedReboot, then wait for it.
	if status.CurrentOperation != updateengine.UpdateStatusUpdatedNeedReboot {
		ctx, cancel := context.WithCancel(context.Background())
		stc := make(chan updateengine.Status)
		go k.ue.RebootNeededSignal(stc, ctx.Done())
		<-stc
		cancel()
	}

	// we are now in UpdateStatusUpdatedNeedReboot state.
	return nil
}

// waitForOkToReboot waits for both 'ok-to-reboot' and 'needs-reboot' to be true.
func (k *Klocksmith) waitForOkToReboot() error {
	n, err := k.nc.Get(k.node, v1meta.GetOptions{})
	if err != nil {
		return fmt.Errorf("failed to get self node (%q): %v", k.node, err)
	}

	if n.Annotations[constants.AnnotationOkToReboot] == constants.True && n.Annotations[constants.AnnotationRebootNeeded] == constants.True {
		return nil
	}

	// XXX: set timeout > 0?
	watcher, err := k.nc.Watch(v1meta.ListOptions{
		FieldSelector:   fields.OneTermEqualSelector("metadata.name", n.Name).String(),
		ResourceVersion: n.ResourceVersion,
	})
	if err != nil {
		return fmt.Errorf("failed to watch self node (%q): %v", k.node, err)
	}

	// hopefully 24 hours is enough time between indicating we need a
	// reboot and the controller telling us to do it
	ev, err := watch.Until(time.Hour*24, watcher, k8sutil.NodeAnnotationCondition(shouldRebootSelector))
	if err != nil {
		return fmt.Errorf("waiting for annotation %q failed: %v", constants.AnnotationOkToReboot, err)
	}

	// sanity check
	no, ok := ev.Object.(*v1.Node)
	if !ok {
		panic("event contains a non-*api.Node object")
	}

	if no.Annotations[constants.AnnotationOkToReboot] != constants.True {
		panic("event did not contain annotation expected")
	}

	return nil
}

func (k *Klocksmith) waitForNotOkToReboot() error {
	n, err := k.nc.Get(k.node, v1meta.GetOptions{})
	if err != nil {
		return fmt.Errorf("failed to get self node (%q): %v", k.node, err)
	}

	if n.Annotations[constants.AnnotationOkToReboot] != constants.True {
		return nil
	}

	// XXX: set timeout > 0?
	watcher, err := k.nc.Watch(v1meta.ListOptions{
		FieldSelector:   fields.OneTermEqualSelector("metadata.name", n.Name).String(),
		ResourceVersion: n.ResourceVersion,
	})
	if err != nil {
		return fmt.Errorf("failed to watch self node (%q): %v", k.node, err)
	}

	// Within 24 hours of indicating we don't need a reboot we should be given a not-ok.
	// If that isn't the case, it likely means the operator isn't running, and
	// we'll just crash-loop in that case, and hopefully that will help the user realize something's wrong.
	// Use a custom condition function to use the more correct 'OkToReboot !=
	// true' vs '== False'; due to the operator matching on '== True', and not
	// going out of its way to convert '' => 'False', checking the exact inverse
	// of what the operator checks is the correct thing to do.
	ev, err := watch.Until(time.Hour*24, watcher, watch.ConditionFunc(func(event watch.Event) (bool, error) {
		switch event.Type {
		case watch.Error:
			return false, fmt.Errorf("error watching node: %v", event.Object)
		case watch.Deleted:
			return false, fmt.Errorf("our node was deleted while we were waiting for ready")
		}

		no := event.Object.(*v1.Node)
		if no.Annotations[constants.AnnotationOkToReboot] != constants.True {
			return true, nil
		}
		return false, nil
	}))
	if err != nil {
		return fmt.Errorf("waiting for annotation %q failed: %v", constants.AnnotationOkToReboot, err)
	}

	// sanity check
	no := ev.Object.(*v1.Node)

	if no.Annotations[constants.AnnotationOkToReboot] == constants.True {
		panic("event did not contain annotation expected")
	}

	return nil
}

func (k *Klocksmith) getPodsForDeletion() ([]v1.Pod, error) {
	pods, err := drain.GetPodsForDeletion(k.kc, k.node)
	if err != nil {
		return nil, fmt.Errorf("failed to get list of pods for deletion: %v", err)
	}

	// XXX: ignoring kube-system is a simple way to avoid eviciting
	// critical components such as kube-scheduler and
	// kube-controller-manager.

	pods = k8sutil.FilterPods(pods, func(p *v1.Pod) bool {
		if p.Namespace == "kube-system" {
			return false
		}

		return true
	})

	return pods, nil
}
<|endoftext|>


<commit_before>package apps

import (
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"time"

	"github.com/cozy/swift"
	"github.com/spf13/afero"
)

// FileServer interface defines a way to access and serve the application's
// data files.
type FileServer interface {
	Open(slug, version, file string) (io.ReadCloser, error)
	ServeFileContent(w http.ResponseWriter, req *http.Request,
		slug, version, file string) error
}

type swiftServer struct {
	c         *swift.Connection
	container string
}

type aferoServer struct {
	mkPath func(slug, version, file string) string
	fs     afero.Fs
}

// NewSwiftFileServer returns provides the apps.FileServer implementation
// using the swift backend as file server.
func NewSwiftFileServer(conn *swift.Connection, appsType AppType) FileServer {
	return &swiftServer{
		c:         conn,
		container: containerName(appsType),
	}
}

func (s *swiftServer) Open(slug, version, file string) (io.ReadCloser, error) {
	objName := s.makeObjectName(slug, version, file)
	f, _, err := s.c.ObjectOpen(s.container, objName, false, nil)
	if err != nil {
		return nil, wrapSwiftErr(err)
	}
	return f, nil
}

func (s *swiftServer) ServeFileContent(w http.ResponseWriter, req *http.Request, slug, version, file string) error {
	objName := s.makeObjectName(slug, version, file)
	f, o, err := s.c.ObjectOpen(s.container, objName, false, nil)
	if err != nil {
		return wrapSwiftErr(err)
	}
	defer f.Close()

	lastModified, _ := time.Parse(http.TimeFormat, o["Last-Modified"])
	if req.Header.Get("Range") == "" {
		w.Header().Set("Etag", fmt.Sprintf(`"%s"`, o["Etag"]))
	}

	http.ServeContent(w, req, objName, lastModified, f)
	return nil
}

func (s *swiftServer) makeObjectName(slug, version, file string) string {
	return path.Join(slug, version, file)
}

// NewAferoFileServer returns a simple wrapper of the afero.Fs interface that
// provides the apps.FileServer interface.
//
// You can provide a makePath method to define how the file name should be
// created from the application's slug, version and file name. If not provided,
// the standard VFS concatenation (starting with vfs.WebappsDirName) is used.
func NewAferoFileServer(fs afero.Fs, makePath func(slug, version, file string) string) FileServer {
	if makePath == nil {
		makePath = defaultMakePath
	}
	return &aferoServer{
		mkPath: makePath,
		fs:     fs,
	}
}

func (s *aferoServer) Open(slug, version, file string) (io.ReadCloser, error) {
	filepath := s.mkPath(slug, version, file)
	f, err := s.open(filepath)
	if os.IsNotExist(err) {
		return s.open(retroCompatMakePath(slug, version, file))
	}
	return f, err
}
func (s *aferoServer) open(filepath string) (io.ReadCloser, error) {
	return s.fs.Open(filepath)
}

func (s *aferoServer) ServeFileContent(w http.ResponseWriter, req *http.Request, slug, version, file string) error {
	filepath := s.mkPath(slug, version, file)
	err := s.serveFileContent(w, req, filepath)
	if os.IsNotExist(err) {
		return s.serveFileContent(w, req, retroCompatMakePath(slug, version, file))
	}
	return err
}
func (s *aferoServer) serveFileContent(w http.ResponseWriter, req *http.Request, filepath string) error {
	infos, err := s.fs.Stat(filepath)
	if err != nil {
		return err
	}

	rc, err := s.fs.Open(filepath)
	if err != nil {
		return err
	}
	defer rc.Close()

	h := md5.New()
	r := io.TeeReader(rc, h)
	b, err := ioutil.ReadAll(r)
	if err != nil {
		return err
	}

	if req.Header.Get("Range") == "" {
		w.Header().Set("Etag", fmt.Sprintf(`"%s"`, hex.EncodeToString(h.Sum(nil))))
	}

	http.ServeContent(w, req, filepath, infos.ModTime(), bytes.NewReader(b))
	return nil
}

func defaultMakePath(slug, version, file string) string {
	basepath := path.Join("/", slug, version)
	filepath := path.Join("/", file)
	return path.Join(basepath, filepath)
}

// FIXME: retro-compatibility code to serve application that were not installed
// in a versioned directory.
func retroCompatMakePath(slug, version, file string) string {
	return path.Join("/", slug, file)
}

func containerName(appsType AppType) string {
	switch appsType {
	case Webapp:
		return "apps-web"
	case Konnector:
		return "apps-konnectors"
	}
	panic("Unknown AppType")
}

func wrapSwiftErr(err error) error {
	if err == swift.ObjectNotFound || err == swift.ContainerNotFound {
		return os.ErrNotExist
	}
	return err
}
<commit_msg>No need to handle last-modified-since<commit_after>package apps

import (
	"bytes"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"time"

	"github.com/cozy/swift"
	"github.com/spf13/afero"
)

var unixZeroEpoch = time.Time{}

// FileServer interface defines a way to access and serve the application's
// data files.
type FileServer interface {
	Open(slug, version, file string) (io.ReadCloser, error)
	ServeFileContent(w http.ResponseWriter, req *http.Request,
		slug, version, file string) error
}

type swiftServer struct {
	c         *swift.Connection
	container string
}

type aferoServer struct {
	mkPath func(slug, version, file string) string
	fs     afero.Fs
}

// NewSwiftFileServer returns provides the apps.FileServer implementation
// using the swift backend as file server.
func NewSwiftFileServer(conn *swift.Connection, appsType AppType) FileServer {
	return &swiftServer{
		c:         conn,
		container: containerName(appsType),
	}
}

func (s *swiftServer) Open(slug, version, file string) (io.ReadCloser, error) {
	objName := s.makeObjectName(slug, version, file)
	f, _, err := s.c.ObjectOpen(s.container, objName, false, nil)
	if err != nil {
		return nil, wrapSwiftErr(err)
	}
	return f, nil
}

func (s *swiftServer) ServeFileContent(w http.ResponseWriter, req *http.Request, slug, version, file string) error {
	objName := s.makeObjectName(slug, version, file)
	f, o, err := s.c.ObjectOpen(s.container, objName, false, nil)
	if err != nil {
		return wrapSwiftErr(err)
	}
	defer f.Close()

	w.Header().Set("Cache-Control", "no-cache")
	if req.Header.Get("Range") == "" {
		w.Header().Set("Etag", fmt.Sprintf(`"%s"`, o["Etag"]))
	}

	http.ServeContent(w, req, objName, unixZeroEpoch, f)
	return nil
}

func (s *swiftServer) makeObjectName(slug, version, file string) string {
	return path.Join(slug, version, file)
}

// NewAferoFileServer returns a simple wrapper of the afero.Fs interface that
// provides the apps.FileServer interface.
//
// You can provide a makePath method to define how the file name should be
// created from the application's slug, version and file name. If not provided,
// the standard VFS concatenation (starting with vfs.WebappsDirName) is used.
func NewAferoFileServer(fs afero.Fs, makePath func(slug, version, file string) string) FileServer {
	if makePath == nil {
		makePath = defaultMakePath
	}
	return &aferoServer{
		mkPath: makePath,
		fs:     fs,
	}
}

func (s *aferoServer) Open(slug, version, file string) (io.ReadCloser, error) {
	filepath := s.mkPath(slug, version, file)
	f, err := s.open(filepath)
	if os.IsNotExist(err) {
		return s.open(retroCompatMakePath(slug, version, file))
	}
	return f, err
}
func (s *aferoServer) open(filepath string) (io.ReadCloser, error) {
	return s.fs.Open(filepath)
}

func (s *aferoServer) ServeFileContent(w http.ResponseWriter, req *http.Request, slug, version, file string) error {
	filepath := s.mkPath(slug, version, file)
	err := s.serveFileContent(w, req, filepath)
	if os.IsNotExist(err) {
		return s.serveFileContent(w, req, retroCompatMakePath(slug, version, file))
	}
	return err
}
func (s *aferoServer) serveFileContent(w http.ResponseWriter, req *http.Request, filepath string) error {
	rc, err := s.fs.Open(filepath)
	if err != nil {
		return err
	}
	defer rc.Close()

	h := md5.New()
	r := io.TeeReader(rc, h)
	b, err := ioutil.ReadAll(r)
	if err != nil {
		return err
	}

	if req.Header.Get("Range") == "" {
		w.Header().Set("Etag", fmt.Sprintf(`"%s"`, hex.EncodeToString(h.Sum(nil))))
	}

	http.ServeContent(w, req, filepath, unixZeroEpoch, bytes.NewReader(b))
	return nil
}

func defaultMakePath(slug, version, file string) string {
	basepath := path.Join("/", slug, version)
	filepath := path.Join("/", file)
	return path.Join(basepath, filepath)
}

// FIXME: retro-compatibility code to serve application that were not installed
// in a versioned directory.
func retroCompatMakePath(slug, version, file string) string {
	return path.Join("/", slug, file)
}

func containerName(appsType AppType) string {
	switch appsType {
	case Webapp:
		return "apps-web"
	case Konnector:
		return "apps-konnectors"
	}
	panic("Unknown AppType")
}

func wrapSwiftErr(err error) error {
	if err == swift.ObjectNotFound || err == swift.ContainerNotFound {
		return os.ErrNotExist
	}
	return err
}
<|endoftext|>


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

package main

import (
	"log"
	"os"
	"path/filepath"

	"github.com/u-root/u-root/pkg/bb"
)

func run() {
	name := filepath.Base(os.Args[0])
	if err := bb.Run(name); err != nil {
		log.Fatalf("%s: %v", name, err)
	}
}

func main() {
	arg1 := os.Args[0]
	for s, err := os.Readlink(arg1); err == nil && s != "bb"; s, err = os.Readlink(arg1) {
		arg1 = s
	}
	os.Args[0] = arg1

	run()
}

func init() {
	m := func() {
		if len(os.Args) == 0 {
			log.Fatal("Arg len is 0. This is impossible")
		}
		if len(os.Args) == 1 {
			// This might be a symlink, and have been invoked by an sshd.
			// Let's try this: readlink until we get a terminal link.
			// If the final link is "", then forget it.
			var arg1 string
			for s, err := os.Readlink(os.Args[0]); err == nil && s != "bb"; s, err = os.Readlink(arg1) {
				arg1 = s
			}
			if arg1 == "" {
				log.Fatalf("os.Args is %v: you need to specify which command to invoke.", os.Args)
			}
			os.Args = append(os.Args, arg1)
		}
		// Use argv[1] as the name.
		os.Args = os.Args[1:]
		run()
	}
	bb.Register("bb", bb.Noop, m)
	bb.RegisterDefault(bb.Noop, m)
}
<commit_msg>bb: check filepath.Base of Args[0]<commit_after>// Copyright 2018 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"log"
	"os"
	"path/filepath"

	"github.com/u-root/u-root/pkg/bb"
)

func run() {
	name := filepath.Base(os.Args[0])
	if err := bb.Run(name); err != nil {
		log.Fatalf("%s: %v", name, err)
	}
}

func main() {
	arg1 := os.Args[0]
	for s, err := os.Readlink(arg1); err == nil && filepath.Base(s) != "bb"; s, err = os.Readlink(arg1) {
		arg1 = s
	}
	os.Args[0] = arg1

	run()
}

func init() {
	m := func() {
		if len(os.Args) == 0 {
			log.Fatal("Arg len is 0. This is impossible")
		}
		if len(os.Args) == 1 {
			// This might be a symlink, and have been invoked by an sshd.
			// Let's try this: readlink until we get a terminal link.
			// If the final link is "", then forget it.
			var arg1 string
			for s, err := os.Readlink(os.Args[0]); err == nil && filepath.Base(s) != "bb"; s, err = os.Readlink(arg1) {
				arg1 = s
			}
			if arg1 == "" {
				log.Fatalf("os.Args is %v: you need to specify which command to invoke.", os.Args)
			}
			os.Args = append(os.Args, arg1)
		}
		// Use argv[1] as the name.
		os.Args = os.Args[1:]
		run()
	}
	bb.Register("bb", bb.Noop, m)
	bb.RegisterDefault(bb.Noop, m)
}
<|endoftext|>


<commit_before>// Copyright Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package envoy

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"istio.io/istio/pkg/http"
	"istio.io/istio/pkg/util/sets"
	"istio.io/pkg/log"
)

var errAbort = errors.New("epoch aborted")

const errOutOfMemory = "signal: killed"

var activeConnectionCheckDelay = 1 * time.Second

// NewAgent creates a new proxy agent for the proxy start-up and clean-up functions.
func NewAgent(proxy Proxy, terminationDrainDuration, minDrainDuration time.Duration, localhost string,
	adminPort, statusPort, prometheusPort int, exitOnZeroActiveConnections bool,
) *Agent {
	knownIstioListeners := sets.New(
		fmt.Sprintf("listener.0.0.0.0_%d.downstream_cx_active", statusPort),
		fmt.Sprintf("listener.0.0.0.0_%d.downstream_cx_active", prometheusPort),
		"listener.admin.downstream_cx_active",
		"listener.admin.main_thread.downstream_cx_active",
	)
	return &Agent{
		proxy:                       proxy,
		statusCh:                    make(chan exitStatus, 1), // context might stop drainage
		drainCh:                     make(chan struct{}),
		abortCh:                     make(chan error, 1),
		terminationDrainDuration:    terminationDrainDuration,
		minDrainDuration:            minDrainDuration,
		exitOnZeroActiveConnections: exitOnZeroActiveConnections,
		adminPort:                   adminPort,
		statusPort:                  statusPort,
		prometheusPort:              prometheusPort,
		localhost:                   localhost,
		knownIstioListeners:         knownIstioListeners,
	}
}

// Proxy defines command interface for a proxy
type Proxy interface {
	// Run command for an epoch, and abort channel
	Run(int, <-chan error) error

	// Drains the current epoch.
	Drain() error

	// Cleanup command for an epoch
	Cleanup(int)

	// UpdateConfig writes a new config file
	UpdateConfig(config []byte) error
}

type Agent struct {
	// proxy commands
	proxy Proxy

	// channel for proxy exit notifications
	statusCh chan exitStatus

	drainCh chan struct{}

	abortCh chan error

	// time to allow for the proxy to drain before terminating all remaining proxy processes
	terminationDrainDuration time.Duration
	minDrainDuration         time.Duration

	adminPort int
	localhost string

	statusPort     int
	prometheusPort int

	knownIstioListeners sets.Set

	exitOnZeroActiveConnections bool
}

type exitStatus struct {
	epoch int
	err   error
}

// Run starts the envoy and waits until it terminates.
func (a *Agent) Run(ctx context.Context) {
	log.Info("Starting proxy agent")
	go a.runWait(0, a.abortCh)

	select {
	case status := <-a.statusCh:
		if status.err != nil {
			if status.err.Error() == errOutOfMemory {
				log.Warnf("Envoy may have been out of memory killed. Check memory usage and limits.")
			}
			log.Errorf("Epoch %d exited with error: %v", status.epoch, status.err)
		} else {
			log.Infof("Epoch %d exited normally", status.epoch)
		}

		log.Infof("No more active epochs, terminating")
	case <-ctx.Done():
		a.terminate()
		status := <-a.statusCh
		if status.err == errAbort {
			log.Infof("Epoch %d aborted normally", status.epoch)
		} else {
			log.Warnf("Epoch %d aborted abnormally", status.epoch)
		}
		log.Info("Agent has successfully terminated")
	}
}

func (a *Agent) terminate() {
	log.Infof("Agent draining Proxy")
	e := a.proxy.Drain()
	if e != nil {
		log.Warnf("Error in invoking drain listeners endpoint %v", e)
	}
	// If exitOnZeroActiveConnections is enabled, always sleep minimumDrainDuration then exit
	// after min(all connections close, terminationGracePeriodSeconds-minimumDrainDuration).
	// exitOnZeroActiveConnections is disabled (default), retain the existing behavior.
	if a.exitOnZeroActiveConnections {
		log.Infof("Agent draining proxy for %v, then waiting for active connections to terminate...", a.minDrainDuration)
		time.Sleep(a.minDrainDuration)
		log.Infof("Checking for active connections...")
		ticker := time.NewTicker(activeConnectionCheckDelay)
		for range ticker.C {
			ac, err := a.activeProxyConnections()
			if err != nil {
				log.Errorf(err.Error())
				a.abortCh <- errAbort
				return
			}
			if ac == -1 {
				log.Info("downstream_cx_active are not available. This either means there are no downstream connection established yet" +
					" or the stats are not enabled. Skipping active connections check...")
				a.abortCh <- errAbort
				return
			}
			if ac == 0 {
				log.Info("There are no more active connections. terminating proxy...")
				a.abortCh <- errAbort
				return
			}
			log.Infof("There are still %d active connections", ac)
		}
	} else {
		log.Infof("Graceful termination period is %v, starting...", a.terminationDrainDuration)
		time.Sleep(a.terminationDrainDuration)
		log.Infof("Graceful termination period complete, terminating remaining proxies.")
		a.abortCh <- errAbort
	}
	log.Warnf("Aborted all epochs")
}

func (a *Agent) activeProxyConnections() (int, error) {
	activeConnectionsURL := fmt.Sprintf("http://%s:%d/stats?usedonly&filter=downstream_cx_active$", a.localhost, a.adminPort)
	stats, err := http.DoHTTPGet(activeConnectionsURL)
	if err != nil {
		return -1, fmt.Errorf("unable to get listener stats from Envoy : %v", err)
	}
	if stats.Len() == 0 {
		return -1, nil
	}
	activeConnections := 0
	for stats.Len() > 0 {
		line, _ := stats.ReadString('\n')
		parts := strings.Split(line, ":")
		if len(parts) != 2 {
			log.Warnf("envoy stat line is missing separator. line:%s", line)
			continue
		}
		// downstream_cx_active is accounted under "http." and "listener." for http listeners.
		// Only consider listener stats. Listener stats also will have per worker stats, we can
		// ignore them.
		if !strings.HasPrefix(parts[0], "listener.") || strings.Contains(parts[0], "worker_") {
			continue
		}
		// If the stat is for a known Istio listener skip it.
		if a.knownIstioListeners.Contains(parts[0]) {
			continue
		}
		val, err := strconv.ParseUint(strings.TrimSpace(parts[1]), 10, 64)
		if err != nil {
			log.Warnf("failed parsing Envoy stat %s (error: %s) line: %s", parts[0], err.Error(), line)
			continue
		}
		activeConnections += int(val)
	}
	if activeConnections > 0 {
		log.Debugf("Active connections stats: %s", stats.String())
	}
	return activeConnections, nil
}

// runWait runs the start-up command as a go routine and waits for it to finish
func (a *Agent) runWait(epoch int, abortCh <-chan error) {
	log.Infof("Epoch %d starting", epoch)
	err := a.proxy.Run(epoch, abortCh)
	a.proxy.Cleanup(epoch)
	a.statusCh <- exitStatus{epoch: epoch, err: err}
}
<commit_msg>Remove unused fields of Agent (#40262)<commit_after>// Copyright Istio Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package envoy

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"istio.io/istio/pkg/http"
	"istio.io/istio/pkg/util/sets"
	"istio.io/pkg/log"
)

var errAbort = errors.New("epoch aborted")

const errOutOfMemory = "signal: killed"

var activeConnectionCheckDelay = 1 * time.Second

// NewAgent creates a new proxy agent for the proxy start-up and clean-up functions.
func NewAgent(proxy Proxy, terminationDrainDuration, minDrainDuration time.Duration, localhost string,
	adminPort, statusPort, prometheusPort int, exitOnZeroActiveConnections bool,
) *Agent {
	knownIstioListeners := sets.New(
		fmt.Sprintf("listener.0.0.0.0_%d.downstream_cx_active", statusPort),
		fmt.Sprintf("listener.0.0.0.0_%d.downstream_cx_active", prometheusPort),
		"listener.admin.downstream_cx_active",
		"listener.admin.main_thread.downstream_cx_active",
	)
	return &Agent{
		proxy:                       proxy,
		statusCh:                    make(chan exitStatus, 1), // context might stop drainage
		abortCh:                     make(chan error, 1),
		terminationDrainDuration:    terminationDrainDuration,
		minDrainDuration:            minDrainDuration,
		exitOnZeroActiveConnections: exitOnZeroActiveConnections,
		adminPort:                   adminPort,
		localhost:                   localhost,
		knownIstioListeners:         knownIstioListeners,
	}
}

// Proxy defines command interface for a proxy
type Proxy interface {
	// Run command for an epoch, and abort channel
	Run(int, <-chan error) error

	// Drains the current epoch.
	Drain() error

	// Cleanup command for an epoch
	Cleanup(int)

	// UpdateConfig writes a new config file
	UpdateConfig(config []byte) error
}

type Agent struct {
	// proxy commands
	proxy Proxy

	// channel for proxy exit notifications
	statusCh chan exitStatus

	abortCh chan error

	// time to allow for the proxy to drain before terminating all remaining proxy processes
	terminationDrainDuration time.Duration
	minDrainDuration         time.Duration

	adminPort int
	localhost string

	knownIstioListeners sets.Set

	exitOnZeroActiveConnections bool
}

type exitStatus struct {
	epoch int
	err   error
}

// Run starts the envoy and waits until it terminates.
func (a *Agent) Run(ctx context.Context) {
	log.Info("Starting proxy agent")
	go a.runWait(0, a.abortCh)

	select {
	case status := <-a.statusCh:
		if status.err != nil {
			if status.err.Error() == errOutOfMemory {
				log.Warnf("Envoy may have been out of memory killed. Check memory usage and limits.")
			}
			log.Errorf("Epoch %d exited with error: %v", status.epoch, status.err)
		} else {
			log.Infof("Epoch %d exited normally", status.epoch)
		}

		log.Infof("No more active epochs, terminating")
	case <-ctx.Done():
		a.terminate()
		status := <-a.statusCh
		if status.err == errAbort {
			log.Infof("Epoch %d aborted normally", status.epoch)
		} else {
			log.Warnf("Epoch %d aborted abnormally", status.epoch)
		}
		log.Info("Agent has successfully terminated")
	}
}

func (a *Agent) terminate() {
	log.Infof("Agent draining Proxy")
	e := a.proxy.Drain()
	if e != nil {
		log.Warnf("Error in invoking drain listeners endpoint %v", e)
	}
	// If exitOnZeroActiveConnections is enabled, always sleep minimumDrainDuration then exit
	// after min(all connections close, terminationGracePeriodSeconds-minimumDrainDuration).
	// exitOnZeroActiveConnections is disabled (default), retain the existing behavior.
	if a.exitOnZeroActiveConnections {
		log.Infof("Agent draining proxy for %v, then waiting for active connections to terminate...", a.minDrainDuration)
		time.Sleep(a.minDrainDuration)
		log.Infof("Checking for active connections...")
		ticker := time.NewTicker(activeConnectionCheckDelay)
		for range ticker.C {
			ac, err := a.activeProxyConnections()
			if err != nil {
				log.Errorf(err.Error())
				a.abortCh <- errAbort
				return
			}
			if ac == -1 {
				log.Info("downstream_cx_active are not available. This either means there are no downstream connection established yet" +
					" or the stats are not enabled. Skipping active connections check...")
				a.abortCh <- errAbort
				return
			}
			if ac == 0 {
				log.Info("There are no more active connections. terminating proxy...")
				a.abortCh <- errAbort
				return
			}
			log.Infof("There are still %d active connections", ac)
		}
	} else {
		log.Infof("Graceful termination period is %v, starting...", a.terminationDrainDuration)
		time.Sleep(a.terminationDrainDuration)
		log.Infof("Graceful termination period complete, terminating remaining proxies.")
		a.abortCh <- errAbort
	}
	log.Warnf("Aborted all epochs")
}

func (a *Agent) activeProxyConnections() (int, error) {
	activeConnectionsURL := fmt.Sprintf("http://%s:%d/stats?usedonly&filter=downstream_cx_active$", a.localhost, a.adminPort)
	stats, err := http.DoHTTPGet(activeConnectionsURL)
	if err != nil {
		return -1, fmt.Errorf("unable to get listener stats from Envoy : %v", err)
	}
	if stats.Len() == 0 {
		return -1, nil
	}
	activeConnections := 0
	for stats.Len() > 0 {
		line, _ := stats.ReadString('\n')
		parts := strings.Split(line, ":")
		if len(parts) != 2 {
			log.Warnf("envoy stat line is missing separator. line:%s", line)
			continue
		}
		// downstream_cx_active is accounted under "http." and "listener." for http listeners.
		// Only consider listener stats. Listener stats also will have per worker stats, we can
		// ignore them.
		if !strings.HasPrefix(parts[0], "listener.") || strings.Contains(parts[0], "worker_") {
			continue
		}
		// If the stat is for a known Istio listener skip it.
		if a.knownIstioListeners.Contains(parts[0]) {
			continue
		}
		val, err := strconv.ParseUint(strings.TrimSpace(parts[1]), 10, 64)
		if err != nil {
			log.Warnf("failed parsing Envoy stat %s (error: %s) line: %s", parts[0], err.Error(), line)
			continue
		}
		activeConnections += int(val)
	}
	if activeConnections > 0 {
		log.Debugf("Active connections stats: %s", stats.String())
	}
	return activeConnections, nil
}

// runWait runs the start-up command as a go routine and waits for it to finish
func (a *Agent) runWait(epoch int, abortCh <-chan error) {
	log.Infof("Epoch %d starting", epoch)
	err := a.proxy.Run(epoch, abortCh)
	a.proxy.Cleanup(epoch)
	a.statusCh <- exitStatus{epoch: epoch, err: err}
}
<|endoftext|>


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

// Package forth implements Forth parsing, which allows
// programs to use forth-like syntax to manipulate a stack
// of Cells.
// It is designed for use by programs
// needing to evaluate command-line arguments or simple
// expressions to set program variables. It is designed
// to map host names to numbers. We use it to
// easily convert host names and IP addresses into
// parameters.
// The language
// is a Forth-like postfix notation. Elements are
// either commands or strings. Strings are
// immediately pushed. Commands consume stack variables
// and produce new ones.
// Simple examples:
// push hostname, strip alpha characters to produce a number. If your
// hostname is sb47, top of stack will be left with 47.
// hostname  hostbase
// Get the hostbase, if it is 0 mod 20, return the hostbase / 20,
// else return hostbase mod 20
//
// hostname hostbase dup 20 / swap 20 % dup ifelse
//
// At the end of the evaluation the stack should have one element
// left; that element is popped and returned. It is an error (currently)
// to return with a non-empty stack.
// This package was used for real work at Sandia National Labs from 2010 to 2012 and possibly later.
// Some of the use of error may seem a bit weird but the creation of this package predates the
// creation of the error type (it was still an os thing back then).
package forth

import (
	"errors"
	"fmt"
	"os"
	"runtime"
	"strconv"
	"strings"
	"sync"
)

type (
	// Op is an opcode type. It does not return an error value,
	// instead using panic when parsing issues occur, to ease
	// the programming annoyance of propagating errors up the
	// stack (following common Go practice for parsers).
	// If you write an op it can use panic as well.
	// Lest you get upset about the use of panic, be aware
	// I've talked to the folks in Go core about this and
	// they feel it's fine for parsers.
	Op func(f Forth)
	// Cell is a stack element.
	Cell interface{}

	stack struct {
		stack []Cell
	}
)

var (
	// Debug is an empty function that can be set to, e.g.,
	// fmt.Printf or log.Printf for debugging.
	Debug   = func(string, ...interface{}) {}
	opmap   map[string]Op
	mapLock sync.Mutex
)

func init() {
	opmap = map[string]Op{
		"+":        plus,
		"-":        sub,
		"*":        times,
		"/":        div,
		"%":        mod,
		"swap":     swap,
		"ifelse":   ifelse,
		"hostname": hostname,
		"hostbase": hostbase,
		"strcat":   strcat,
		"roundup":  roundup,
		"dup":      dup,
		"drop":     drop,
		"newword":  newword,
		"words":    words,
	}
}

// Forth is an interface used by the package. The interface
// requires definition of Push, Pop, Length, Empty (convenience function
// meaning Length is 0), Newop (insert a new or replacement operator),
// and Reset (clear the stack, mainly diagnostic)
type Forth interface {
	Push(Cell)
	Pop() Cell
	Length() int
	Empty() bool
	Reset()
	Stack() []Cell
}

// New creates a new stack
func New() Forth {
	f := new(stack)
	return f
}

// Getop gets an op from the map.
func Getop(n string) Op {
	mapLock.Lock()
	defer mapLock.Unlock()
	op, ok := opmap[n]
	if !ok {
		return nil
	}
	return op
}

// Putop creates a new operation. We considered having
// an opmap per stack but don't feel the package requires it
func Putop(n string, op Op) {
	mapLock.Lock()
	defer mapLock.Unlock()
	if _, ok := opmap[n]; ok {
		panic("Putting %s: op already assigned")
	}
	opmap[n] = op
}

// Ops returns the operator map.
func Ops() map[string]Op {
	return opmap
}

// Reset resets the stack to empty
func (f *stack) Reset() {
	f.stack = f.stack[0:0]
}

// Return the stack
func (f *stack) Stack() []Cell {
	return f.stack
}

// Push pushes the interface{} on the stack.
func (f *stack) Push(c Cell) {
	f.stack = append(f.stack, c)
	Debug("push: %v: stack: %v\n", c, f.stack)
}

// Pop pops the stack. If the stack is Empty Pop will panic.
// Eval recovers() the panic.
func (f *stack) Pop() (ret Cell) {
	if len(f.stack) < 1 {
		panic(errors.New("Empty stack"))
	}
	ret = f.stack[len(f.stack)-1]
	f.stack = f.stack[0 : len(f.stack)-1]
	Debug("Pop: %v stack %v\n", ret, f.stack)
	return ret
}

// Length returns the stack length.
func (f *stack) Length() int {
	return len(f.stack)
}

// Empty is a convenience function synonymous with Length == 0
func (f *stack) Empty() bool {
	return len(f.stack) == 0
}

// errRecover converts panics to errstr iff it is an os.Error, panics
// otherwise.
func errRecover(errp *error) {
	e := recover()
	if e != nil {
		if _, ok := e.(runtime.Error); ok {
			Debug("errRecover panics with a runtime error")
			panic(e)
		}
		Debug("errRecover returns %v", e)
		*errp = fmt.Errorf("%v", e)
	}
}

// Eval takes a Forth and []Cell, pushing each element on the stack or invoking the
// operator if it is found in the opmap.
func eval(f Forth, cells ...Cell) {
	Debug("eval cells %v", cells)
	for _, c := range cells {
		Debug("eval %v(%T) stack %v", c, c, f.Stack())
		switch s := c.(type) {
		case string:
			fun := Getop(s)
			if fun != nil {
				Debug("eval ... %v:", f.Stack())
				fun(f)
				Debug("eval: Stack now %v", f.Stack())
				break
			}
			if s[0] == '\'' {
				s = s[1:]
			}
			f.Push(s)
			Debug("push %v(%T), stack %v", s, s, f.Stack())
		default:
			Debug("push %v(%T), stack %v", s, s, f.Stack())
			f.Push(s)
		}
	}
}

// Eval calls eval and catches panics.
func Eval(f Forth, cells ...Cell) (err error) {
	defer errRecover(&err)
	eval(f, cells...)
	return
}

// EvalString takes a Forth and string and splits the string on space
// characters, calling Eval for each one.
func EvalString(f Forth, s string) (err error) {
	for _, c := range strings.Fields(s) {
		if err = Eval(f, c); err != nil {
			return
		}
	}
	Debug("EvalString err %v", err)
	return
}

// EvalPop takes a Forth and string, calls EvalString, and
// returns TOS and an error, if any.
// For EvalPop it is an error to leave the stack non-Empty.
// EvalPop is typically used for programs that want to
// parse forth contained in, e.g., flag.Args(), and return
// a result. In most use cases, we want the stack to be empty.
func EvalPop(f Forth, s string) (ret Cell, err error) {
	defer errRecover(&err)
	if err = EvalString(f, s); err != nil {
		return
	}
	if f.Length() != 1 {
		panic(fmt.Sprintf("%v: length is not 1", f.Stack()))
	}
	ret = f.Pop()
	Debug("EvalPop ret %v err %v", ret, err)
	return
}

// String returns the Top Of Stack if it is a string
// or panics.
func String(f Forth) string {
	c := f.Pop()
	switch s := c.(type) {
	case string:
		return s
	default:
		panic(fmt.Sprintf("Can't convert %v to a string", c))
	}
}

// toInt converts to int64.
func toInt(f Forth) int64 {
	Debug("toint %v", f.Stack())
	c := f.Pop()
	Debug("%T", c)
	switch s := c.(type) {
	case string:
		i, err := strconv.ParseInt(s, 0, 64)
		if err != nil {
			panic(err)
		}
		return i
	case int64:
		return s
	default:
		panic(fmt.Sprintf("NaN: %T", c))
	}
}

func plus(f Forth) {
	x := toInt(f)
	y := toInt(f)
	z := x + y
	f.Push(strconv.FormatInt(z, 10))
}

func words(f Forth) {
	mapLock.Lock()
	defer mapLock.Unlock()
	var w []string
	for i := range opmap {
		w = append(w, i)
	}
	f.Push(w)
}

func newword(f Forth) {
	s := String(f)
	n := toInt(f)
	// Pop <n> Cells.
	if f.Length() < int(n) {
		panic(fmt.Sprintf("newword %s: stack is %d elements, need %d", s, f.Length(), n))
	}
	c := make([]Cell, n)
	for i := n; i > 0; i-- {
		c[i-1] = f.Pop()
	}
	Putop(s, func(f Forth) {
		Debug("c %v", c)
		eval(f, c...)
	})
}

func drop(f Forth) {
	_ = f.Pop()
}

func times(f Forth) {
	x := toInt(f)
	y := toInt(f)
	z := x * y
	f.Push(strconv.FormatInt(z, 10))
}

func sub(f Forth) {
	x := toInt(f)
	y := toInt(f)
	z := y - x
	f.Push(strconv.FormatInt(z, 10))
}

func div(f Forth) {
	x := toInt(f)
	y := toInt(f)
	z := y / x
	f.Push(strconv.FormatInt(z, 10))
}

func mod(f Forth) {
	x := toInt(f)
	y := toInt(f)
	z := y % x
	f.Push(strconv.FormatInt(z, 10))
}

func roundup(f Forth) {
	rnd := toInt(f)
	v := toInt(f)
	v = ((v + rnd - 1) / rnd) * rnd
	f.Push(strconv.FormatInt(v, 10))
}

func swap(f Forth) {
	x := f.Pop()
	y := f.Pop()
	f.Push(x)
	f.Push(y)
}

func strcat(f Forth) {
	x := String(f)
	y := String(f)
	f.Push(y + x)
}

func dup(f Forth) {
	x := f.Pop()
	f.Push(x)
	f.Push(x)
}

func ifelse(f Forth) {
	x := toInt(f)
	y := f.Pop()
	z := f.Pop()
	if x != 0 {
		f.Push(y)
	} else {
		f.Push(z)
	}
}

func hostname(f Forth) {
	h, err := os.Hostname()
	if err != nil {
		panic("No hostname")
	}
	f.Push(h)
}

func hostbase(f Forth) {
	host := String(f)
	f.Push(strings.TrimLeft(host, "abcdefghijklmnopqrstuvwxyz -"))
}

// NewWord allows for definition of new operators from strings.
// For example, should you wish to create a word which adds a number
// to itself twice, you can call:
// NewWord(f, "d3d", "dup dup + +")
// which does two dups, and two adds.
func NewWord(f Forth, name string, cell Cell, cells ...Cell) {
	cmd := append([]Cell{cell}, cells...)
	newword := func(f Forth) {
		eval(f, cmd...)
	}
	Putop(name, newword)
}
<commit_msg>pkg/forth: change casting in stack size test for codequal<commit_after>// Copyright 2010-2018 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package forth implements Forth parsing, which allows
// programs to use forth-like syntax to manipulate a stack
// of Cells.
// It is designed for use by programs
// needing to evaluate command-line arguments or simple
// expressions to set program variables. It is designed
// to map host names to numbers. We use it to
// easily convert host names and IP addresses into
// parameters.
// The language
// is a Forth-like postfix notation. Elements are
// either commands or strings. Strings are
// immediately pushed. Commands consume stack variables
// and produce new ones.
// Simple examples:
// push hostname, strip alpha characters to produce a number. If your
// hostname is sb47, top of stack will be left with 47.
// hostname  hostbase
// Get the hostbase, if it is 0 mod 20, return the hostbase / 20,
// else return hostbase mod 20
//
// hostname hostbase dup 20 / swap 20 % dup ifelse
//
// At the end of the evaluation the stack should have one element
// left; that element is popped and returned. It is an error (currently)
// to return with a non-empty stack.
// This package was used for real work at Sandia National Labs from 2010 to 2012 and possibly later.
// Some of the use of error may seem a bit weird but the creation of this package predates the
// creation of the error type (it was still an os thing back then).
package forth

import (
	"errors"
	"fmt"
	"os"
	"runtime"
	"strconv"
	"strings"
	"sync"
)

type (
	// Op is an opcode type. It does not return an error value,
	// instead using panic when parsing issues occur, to ease
	// the programming annoyance of propagating errors up the
	// stack (following common Go practice for parsers).
	// If you write an op it can use panic as well.
	// Lest you get upset about the use of panic, be aware
	// I've talked to the folks in Go core about this and
	// they feel it's fine for parsers.
	Op func(f Forth)
	// Cell is a stack element.
	Cell interface{}

	stack struct {
		stack []Cell
	}
)

var (
	// Debug is an empty function that can be set to, e.g.,
	// fmt.Printf or log.Printf for debugging.
	Debug   = func(string, ...interface{}) {}
	opmap   map[string]Op
	mapLock sync.Mutex
)

func init() {
	opmap = map[string]Op{
		"+":        plus,
		"-":        sub,
		"*":        times,
		"/":        div,
		"%":        mod,
		"swap":     swap,
		"ifelse":   ifelse,
		"hostname": hostname,
		"hostbase": hostbase,
		"strcat":   strcat,
		"roundup":  roundup,
		"dup":      dup,
		"drop":     drop,
		"newword":  newword,
		"words":    words,
	}
}

// Forth is an interface used by the package. The interface
// requires definition of Push, Pop, Length, Empty (convenience function
// meaning Length is 0), Newop (insert a new or replacement operator),
// and Reset (clear the stack, mainly diagnostic)
type Forth interface {
	Push(Cell)
	Pop() Cell
	Length() int
	Empty() bool
	Reset()
	Stack() []Cell
}

// New creates a new stack
func New() Forth {
	f := new(stack)
	return f
}

// Getop gets an op from the map.
func Getop(n string) Op {
	mapLock.Lock()
	defer mapLock.Unlock()
	op, ok := opmap[n]
	if !ok {
		return nil
	}
	return op
}

// Putop creates a new operation. We considered having
// an opmap per stack but don't feel the package requires it
func Putop(n string, op Op) {
	mapLock.Lock()
	defer mapLock.Unlock()
	if _, ok := opmap[n]; ok {
		panic("Putting %s: op already assigned")
	}
	opmap[n] = op
}

// Ops returns the operator map.
func Ops() map[string]Op {
	return opmap
}

// Reset resets the stack to empty
func (f *stack) Reset() {
	f.stack = f.stack[0:0]
}

// Return the stack
func (f *stack) Stack() []Cell {
	return f.stack
}

// Push pushes the interface{} on the stack.
func (f *stack) Push(c Cell) {
	f.stack = append(f.stack, c)
	Debug("push: %v: stack: %v\n", c, f.stack)
}

// Pop pops the stack. If the stack is Empty Pop will panic.
// Eval recovers() the panic.
func (f *stack) Pop() (ret Cell) {
	if len(f.stack) < 1 {
		panic(errors.New("Empty stack"))
	}
	ret = f.stack[len(f.stack)-1]
	f.stack = f.stack[0 : len(f.stack)-1]
	Debug("Pop: %v stack %v\n", ret, f.stack)
	return ret
}

// Length returns the stack length.
func (f *stack) Length() int {
	return len(f.stack)
}

// Empty is a convenience function synonymous with Length == 0
func (f *stack) Empty() bool {
	return len(f.stack) == 0
}

// errRecover converts panics to errstr iff it is an os.Error, panics
// otherwise.
func errRecover(errp *error) {
	e := recover()
	if e != nil {
		if _, ok := e.(runtime.Error); ok {
			Debug("errRecover panics with a runtime error")
			panic(e)
		}
		Debug("errRecover returns %v", e)
		*errp = fmt.Errorf("%v", e)
	}
}

// Eval takes a Forth and []Cell, pushing each element on the stack or invoking the
// operator if it is found in the opmap.
func eval(f Forth, cells ...Cell) {
	Debug("eval cells %v", cells)
	for _, c := range cells {
		Debug("eval %v(%T) stack %v", c, c, f.Stack())
		switch s := c.(type) {
		case string:
			fun := Getop(s)
			if fun != nil {
				Debug("eval ... %v:", f.Stack())
				fun(f)
				Debug("eval: Stack now %v", f.Stack())
				break
			}
			if s[0] == '\'' {
				s = s[1:]
			}
			f.Push(s)
			Debug("push %v(%T), stack %v", s, s, f.Stack())
		default:
			Debug("push %v(%T), stack %v", s, s, f.Stack())
			f.Push(s)
		}
	}
}

// Eval calls eval and catches panics.
func Eval(f Forth, cells ...Cell) (err error) {
	defer errRecover(&err)
	eval(f, cells...)
	return
}

// EvalString takes a Forth and string and splits the string on space
// characters, calling Eval for each one.
func EvalString(f Forth, s string) (err error) {
	for _, c := range strings.Fields(s) {
		if err = Eval(f, c); err != nil {
			return
		}
	}
	Debug("EvalString err %v", err)
	return
}

// EvalPop takes a Forth and string, calls EvalString, and
// returns TOS and an error, if any.
// For EvalPop it is an error to leave the stack non-Empty.
// EvalPop is typically used for programs that want to
// parse forth contained in, e.g., flag.Args(), and return
// a result. In most use cases, we want the stack to be empty.
func EvalPop(f Forth, s string) (ret Cell, err error) {
	defer errRecover(&err)
	if err = EvalString(f, s); err != nil {
		return
	}
	if f.Length() != 1 {
		panic(fmt.Sprintf("%v: length is not 1", f.Stack()))
	}
	ret = f.Pop()
	Debug("EvalPop ret %v err %v", ret, err)
	return
}

// String returns the Top Of Stack if it is a string
// or panics.
func String(f Forth) string {
	c := f.Pop()
	switch s := c.(type) {
	case string:
		return s
	default:
		panic(fmt.Sprintf("Can't convert %v to a string", c))
	}
}

// toInt converts to int64.
func toInt(f Forth) int64 {
	Debug("toint %v", f.Stack())
	c := f.Pop()
	Debug("%T", c)
	switch s := c.(type) {
	case string:
		i, err := strconv.ParseInt(s, 0, 64)
		if err != nil {
			panic(err)
		}
		return i
	case int64:
		return s
	default:
		panic(fmt.Sprintf("NaN: %T", c))
	}
}

func plus(f Forth) {
	x := toInt(f)
	y := toInt(f)
	z := x + y
	f.Push(strconv.FormatInt(z, 10))
}

func words(f Forth) {
	mapLock.Lock()
	defer mapLock.Unlock()
	var w []string
	for i := range opmap {
		w = append(w, i)
	}
	f.Push(w)
}

func newword(f Forth) {
	s := String(f)
	n := toInt(f)
	// Pop <n> Cells.
	if int64(f.Length()) < n {
		panic(fmt.Sprintf("newword %s: stack is %d elements, need %d", s, f.Length(), n))
	}
	c := make([]Cell, n)
	for i := n; i > 0; i-- {
		c[i-1] = f.Pop()
	}
	Putop(s, func(f Forth) {
		Debug("c %v", c)
		eval(f, c...)
	})
}

func drop(f Forth) {
	_ = f.Pop()
}

func times(f Forth) {
	x := toInt(f)
	y := toInt(f)
	z := x * y
	f.Push(strconv.FormatInt(z, 10))
}

func sub(f Forth) {
	x := toInt(f)
	y := toInt(f)
	z := y - x
	f.Push(strconv.FormatInt(z, 10))
}

func div(f Forth) {
	x := toInt(f)
	y := toInt(f)
	z := y / x
	f.Push(strconv.FormatInt(z, 10))
}

func mod(f Forth) {
	x := toInt(f)
	y := toInt(f)
	z := y % x
	f.Push(strconv.FormatInt(z, 10))
}

func roundup(f Forth) {
	rnd := toInt(f)
	v := toInt(f)
	v = ((v + rnd - 1) / rnd) * rnd
	f.Push(strconv.FormatInt(v, 10))
}

func swap(f Forth) {
	x := f.Pop()
	y := f.Pop()
	f.Push(x)
	f.Push(y)
}

func strcat(f Forth) {
	x := String(f)
	y := String(f)
	f.Push(y + x)
}

func dup(f Forth) {
	x := f.Pop()
	f.Push(x)
	f.Push(x)
}

func ifelse(f Forth) {
	x := toInt(f)
	y := f.Pop()
	z := f.Pop()
	if x != 0 {
		f.Push(y)
	} else {
		f.Push(z)
	}
}

func hostname(f Forth) {
	h, err := os.Hostname()
	if err != nil {
		panic("No hostname")
	}
	f.Push(h)
}

func hostbase(f Forth) {
	host := String(f)
	f.Push(strings.TrimLeft(host, "abcdefghijklmnopqrstuvwxyz -"))
}

// NewWord allows for definition of new operators from strings.
// For example, should you wish to create a word which adds a number
// to itself twice, you can call:
// NewWord(f, "d3d", "dup dup + +")
// which does two dups, and two adds.
func NewWord(f Forth, name string, cell Cell, cells ...Cell) {
	cmd := append([]Cell{cell}, cells...)
	newword := func(f Forth) {
		eval(f, cmd...)
	}
	Putop(name, newword)
}
<|endoftext|>