<commit_before>//
// Blackblog
// Copyright 2012 Google Inc. All rights reserved.
//
// 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 main

import (
	"errors"
	"flag"
	"fmt"
	"net/http"
	"strings"
	"sync"
	"time"
)

var (
	serverPort = flag.Int("port", 0, "The port on which the standalone HTTP server will run.")

	serverPollWait = flag.Int("server-poll-time", 60, "The time in seconds that the server waits before polling the directory for changes.")
)

type blogServer struct {
	root string // Path to the root of the blog.

	mu    *sync.RWMutex
	posts PostList
	r     *render
}

// RunAsServer checks if the program has been configured to run as a web server.
func RunAsServer() bool {
	return *serverPort != 0
}

// StartBlogServer runs the program's web server given the blog located
// at |blogRoot|.
func StartBlogServer(blogRoot string) error {
	if !RunAsServer() {
		return errors.New("No --port specified to start the server")
	}

	server := &blogServer{
		root: blogRoot,
		mu:   new(sync.RWMutex),
	}

	err := server.buildPosts()
	if err != nil {
		return err
	}
	go server.pollPostChanges()

	fmt.Printf("Starting blog server on port %d\n", *serverPort)
	return http.ListenAndServe(fmt.Sprintf(":%d", *serverPort), server)
}

func (b *blogServer) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	url := strings.Trim(req.URL.Path, "/")

	b.mu.RLock()
	defer b.mu.RUnlock()

	if url == "" {
		b.serveNode(rw, req, b.r)
		return
	}

	parts := strings.Split(url, "/")
	node := b.r
	for _, part := range parts {
		if child, ok := node.object.(renderTree)[part]; ok {
			node = child
		} else {
			http.NotFound(rw, req)
			return
		}
	}

	b.serveNode(rw, req, node)
}

func (b *blogServer) serveNode(rw http.ResponseWriter, req *http.Request, render *render) {
	switch render.t {
	case renderTypePost:
		post := render.object.(*Post)
		data, err := post.GetContents()
		if err != nil {
			rw.WriteHeader(http.StatusNotFound)
			fmt.Fprint(rw, err.Error())
			return
		}
		content := RenderPost(post, data)
		rw.Write(content)
	case renderTypeRedirect:
		fallthrough
	case renderTypeDirectory:
		// The root element should generate a post list.
		if render.t == renderTypeDirectory && render.parent == nil {
			index, err := CreateIndex(b.posts)
			if err != nil {
				rw.WriteHeader(http.StatusInternalServerError)
				fmt.Fprint(rw, err.Error())
				return
			}
			rw.Write(index)
			return
		}

		// Other directories when accessed directly should fallback to the
		// redirect.
		if render.t == renderTypeDirectory {
			render = render.object.(renderTree)["index.html"]
		}

		http.Redirect(rw, req, render.object.(string), http.StatusMovedPermanently)
	default:
		rw.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(rw, "Unknown render: %v", render)
	}
}

func (b *blogServer) pollPostChanges() {
	for {
		time.Sleep(time.Duration(*serverPollWait) * time.Second)
		if err := b.buildPosts(); err != nil {
			panic(err.Error())
		}
	}
}

func (b *blogServer) buildPosts() (err error) {
	newPosts := GetPostsInDirectory(b.root)

	b.mu.RLock()
	rebuild := len(newPosts) != len(b.posts)
	if !rebuild {
		for _, p := range b.posts {
			if !p.IsUpToDate() {
				rebuild = true
				break
			}
		}
	}
	b.mu.RUnlock()

	if rebuild {
		b.mu.Lock()
		defer b.mu.Unlock()

		b.posts = GetPostsInDirectory(b.root)
		b.r, err = createRenderTree(b.posts)
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>Fix compile after merge<commit_after>//
// Blackblog
// Copyright 2012 Google Inc. All rights reserved.
//
// 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 main

import (
	"errors"
	"flag"
	"fmt"
	"net/http"
	"strings"
	"sync"
	"time"
)

var (
	serverPort = flag.Int("port", 0, "The port on which the standalone HTTP server will run.")

	serverPollWait = flag.Int("server-poll-time", 60, "The time in seconds that the server waits before polling the directory for changes.")
)

type blogServer struct {
	root string // Path to the root of the blog.

	mu    *sync.RWMutex
	posts PostList
	r     *render
}

// RunAsServer checks if the program has been configured to run as a web server.
func RunAsServer() bool {
	return *serverPort != 0
}

// StartBlogServer runs the program's web server given the blog located
// at |blogRoot|.
func StartBlogServer(blogRoot string) error {
	if !RunAsServer() {
		return errors.New("No --port specified to start the server")
	}

	server := &blogServer{
		root: blogRoot,
		mu:   new(sync.RWMutex),
	}

	err := server.buildPosts()
	if err != nil {
		return err
	}
	go server.pollPostChanges()

	fmt.Printf("Starting blog server on port %d\n", *serverPort)
	return http.ListenAndServe(fmt.Sprintf(":%d", *serverPort), server)
}

func (b *blogServer) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	url := strings.Trim(req.URL.Path, "/")

	b.mu.RLock()
	defer b.mu.RUnlock()

	if url == "" {
		b.serveNode(rw, req, b.r)
		return
	}

	parts := strings.Split(url, "/")
	node := b.r
	for _, part := range parts {
		if child, ok := node.object.(renderTree)[part]; ok {
			node = child
		} else {
			http.NotFound(rw, req)
			return
		}
	}

	b.serveNode(rw, req, node)
}

func (b *blogServer) serveNode(rw http.ResponseWriter, req *http.Request, render *render) {
	switch render.t {
	case renderTypePost:
		post := render.object.(*Post)
		data, err := post.GetContents()
		if err != nil {
			rw.WriteHeader(http.StatusNotFound)
			fmt.Fprint(rw, err.Error())
			return
		}
		content := RenderPost(post, data)
		rw.Write(content)
	case renderTypeRedirect:
		fallthrough
	case renderTypeDirectory:
		// The root element should generate a post list.
		if render.t == renderTypeDirectory && render.parent == nil {
			index, err := CreateIndex(b.posts)
			if err != nil {
				rw.WriteHeader(http.StatusInternalServerError)
				fmt.Fprint(rw, err.Error())
				return
			}
			rw.Write(index)
			return
		}

		// Other directories when accessed directly should fallback to the
		// redirect.
		if render.t == renderTypeDirectory {
			render = render.object.(renderTree)["index.html"]
		}

		http.Redirect(rw, req, render.object.(string), http.StatusMovedPermanently)
	default:
		rw.WriteHeader(http.StatusInternalServerError)
		fmt.Fprintf(rw, "Unknown render: %v", render)
	}
}

func (b *blogServer) pollPostChanges() {
	for {
		time.Sleep(time.Duration(*serverPollWait) * time.Second)
		if err := b.buildPosts(); err != nil {
			panic(err.Error())
		}
	}
}

func (b *blogServer) buildPosts() (err error) {
	newPosts, err := GetPostsInDirectory(b.root)
	if err != nil {
		return
	}

	b.mu.RLock()
	rebuild := len(newPosts) != len(b.posts)
	if !rebuild {
		for _, p := range b.posts {
			if !p.IsUpToDate() {
				rebuild = true
				break
			}
		}
	}
	b.mu.RUnlock()

	if rebuild {
		b.mu.Lock()
		defer b.mu.Unlock()

		b.posts, err = GetPostsInDirectory(b.root)
		if err != nil {
			return
		}
		b.r, err = createRenderTree(b.posts)
		if err != nil {
			return
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/byuoitav/device-monitoring-microservice/statusinfrastructure"
	"github.com/byuoitav/event-router-microservice/eventinfrastructure"
	"github.com/byuoitav/touchpanel-ui-microservice/events"
	"github.com/byuoitav/touchpanel-ui-microservice/handlers"
	"github.com/byuoitav/touchpanel-ui-microservice/socket"
	"github.com/byuoitav/touchpanel-ui-microservice/uiconfig"
	"github.com/jessemillar/health"
	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
)

func main() {
	filters := []string{eventinfrastructure.UI, eventinfrastructure.UIFeature}
	en := eventinfrastructure.NewEventNode("Touchpanel UI", filters, os.Getenv("EVENT_ROUTER_ADDRESS"))

	// websocket hub
	hub := socket.NewHub()
	go events.WriteEventsToSocket(en, hub)
	go events.SendRefresh(hub, time.NewTimer(time.Second*10))

	port := ":8888"
	router := echo.New()
	router.Pre(middleware.RemoveTrailingSlash())
	router.Use(middleware.CORS())

	router.GET("/health", echo.WrapHandler(http.HandlerFunc(health.Check)))
	router.GET("/mstatus", GetStatus)

	// event endpoints
	router.POST("/publish", handlers.PublishEvent, BindEventNode(en))
	router.POST("/publishfeature", handlers.PublishFeature, BindEventNode(en))

	// websocket
	router.GET("/websocket", func(context echo.Context) error {
		socket.ServeWebsocket(hub, context.Response().Writer, context.Request())
		return nil
	})

	// socket endpoints
	router.PUT("/screenoff", func(context echo.Context) error {
		events.SendScreenTimeout(hub)
		return nil
	})
	router.PUT("/refresh", func(context echo.Context) error {
		events.SendRefresh(hub, time.NewTimer(0))
		return nil
	})
	router.GET("/wsinfo", func(context echo.Context) error {
		si, _ := socket.GetSocketInfo(hub)
		return context.JSON(http.StatusOK, si)
	})
	router.PUT("/socketTest", func(context echo.Context) error {
		events.SendTest(hub)
		return context.JSON(http.StatusOK, "sent")
	})

	router.GET("/pihostname", handlers.GetPiHostname)
	router.GET("/hostname", handlers.GetHostname)
	router.GET("/deviceinfo", handlers.GetDeviceInfo)
	router.GET("/reboot", handlers.Reboot)
	router.GET("/dockerstatus", handlers.GetDockerStatus)

	router.GET("/uiconfig", uiconfig.GetUIConfig)
	router.GET("/uipath", uiconfig.GetUIPath)
	router.GET("/api", uiconfig.GetAPI)
	router.GET("/nextapi", uiconfig.NextAPI)

	router.POST("/help", handlers.Help)
	router.POST("/confirmhelp", handlers.ConfirmHelp)
	router.POST("/cancelhelp", handlers.CancelHelp)

	// all the different ui's
	router.Static("/", "redirect.html")
	router.Any("/404", redirect)
	router.Static("/blueberry", "blueberry-dist")

	router.Start(port)
}

func BindEventNode(en *eventinfrastructure.EventNode) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			c.Set(eventinfrastructure.ContextEventNode, en)
			return next(c)
		}
	}
}

func GetStatus(context echo.Context) error {
	var s statusinfrastructure.Status
	var err error

	s.Version, err = statusinfrastructure.GetVersion("version.txt")
	if err != nil {
		s.Version = "missing"
		s.Status = statusinfrastructure.StatusSick
		s.StatusInfo = fmt.Sprintf("Error: %s", err.Error())
	} else {
		s.Status = statusinfrastructure.StatusOK
		s.StatusInfo = ""
	}

	return context.JSON(http.StatusOK, s)
}

func redirect(context echo.Context) error {
	http.Redirect(context.Response().Writer, context.Request(), "http://github.com/404", 302)
	return nil
}
<commit_msg>test authentication<commit_after>package main

import (
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/byuoitav/authmiddleware"
	"github.com/byuoitav/device-monitoring-microservice/statusinfrastructure"
	"github.com/byuoitav/event-router-microservice/eventinfrastructure"
	"github.com/byuoitav/touchpanel-ui-microservice/events"
	"github.com/byuoitav/touchpanel-ui-microservice/handlers"
	"github.com/byuoitav/touchpanel-ui-microservice/socket"
	"github.com/byuoitav/touchpanel-ui-microservice/uiconfig"
	"github.com/jessemillar/health"
	"github.com/labstack/echo"
	"github.com/labstack/echo/middleware"
)

func main() {
	filters := []string{eventinfrastructure.UI, eventinfrastructure.UIFeature}
	en := eventinfrastructure.NewEventNode("Touchpanel UI", filters, os.Getenv("EVENT_ROUTER_ADDRESS"))

	// websocket hub
	hub := socket.NewHub()
	go events.WriteEventsToSocket(en, hub)
	go events.SendRefresh(hub, time.NewTimer(time.Second*10))

	port := ":8888"
	router := echo.New()
	router.Pre(middleware.RemoveTrailingSlash())
	router.Use(middleware.CORS())
	router.Use(echo.WrapMiddleware(authmiddleware.AuthenticateUser))

	router.GET("/health", echo.WrapHandler(http.HandlerFunc(health.Check)))
	router.GET("/mstatus", GetStatus)

	// event endpoints
	router.POST("/publish", handlers.PublishEvent, BindEventNode(en))
	router.POST("/publishfeature", handlers.PublishFeature, BindEventNode(en))

	// websocket
	router.GET("/websocket", func(context echo.Context) error {
		socket.ServeWebsocket(hub, context.Response().Writer, context.Request())
		return nil
	})

	// socket endpoints
	router.PUT("/screenoff", func(context echo.Context) error {
		events.SendScreenTimeout(hub)
		return nil
	})
	router.PUT("/refresh", func(context echo.Context) error {
		events.SendRefresh(hub, time.NewTimer(0))
		return nil
	})
	router.GET("/wsinfo", func(context echo.Context) error {
		si, _ := socket.GetSocketInfo(hub)
		return context.JSON(http.StatusOK, si)
	})
	router.PUT("/socketTest", func(context echo.Context) error {
		events.SendTest(hub)
		return context.JSON(http.StatusOK, "sent")
	})

	router.GET("/pihostname", handlers.GetPiHostname)
	router.GET("/hostname", handlers.GetHostname)
	router.GET("/deviceinfo", handlers.GetDeviceInfo)
	router.GET("/reboot", handlers.Reboot)
	router.GET("/dockerstatus", handlers.GetDockerStatus)

	router.GET("/uiconfig", uiconfig.GetUIConfig)
	router.GET("/uipath", uiconfig.GetUIPath)
	router.GET("/api", uiconfig.GetAPI)
	router.GET("/nextapi", uiconfig.NextAPI)

	router.POST("/help", handlers.Help)
	router.POST("/confirmhelp", handlers.ConfirmHelp)
	router.POST("/cancelhelp", handlers.CancelHelp)

	// all the different ui's
	router.Static("/", "redirect.html")
	router.Any("/404", redirect)
	router.Static("/blueberry", "blueberry-dist")

	router.Start(port)
}

func BindEventNode(en *eventinfrastructure.EventNode) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			c.Set(eventinfrastructure.ContextEventNode, en)
			return next(c)
		}
	}
}

func GetStatus(context echo.Context) error {
	var s statusinfrastructure.Status
	var err error

	s.Version, err = statusinfrastructure.GetVersion("version.txt")
	if err != nil {
		s.Version = "missing"
		s.Status = statusinfrastructure.StatusSick
		s.StatusInfo = fmt.Sprintf("Error: %s", err.Error())
	} else {
		s.Status = statusinfrastructure.StatusOK
		s.StatusInfo = ""
	}

	return context.JSON(http.StatusOK, s)
}

func redirect(context echo.Context) error {
	http.Redirect(context.Response().Writer, context.Request(), "http://github.com/404", 302)
	return nil
}
<|endoftext|>


<commit_before>/*
Server implements RFC 6143 §7.6 Server-to-Client Messages.

See http://tools.ietf.org/html/rfc6143#section-7.6 for more info.
*/
package vnc

import (
	"encoding/binary"
	"fmt"
	"io"
)

const (
	framebufferUpdateMsg = iota
	setColorMapEntriesMsg
	bellMsg
	serverCutTextMsg
)

// A ServerMessage implements a message sent from the server to the client.
type ServerMessage interface {
	// The type of the message that is sent down on the wire.
	Type() uint8

	// Read reads the contents of the message from the reader. At the point
	// this is called, the message type has already been read from the reader.
	// This should return a new ServerMessage that is the appropriate type.
	Read(*ClientConn, io.Reader) (ServerMessage, error)
}

// FramebufferUpdate consists of a sequence of rectangles of
// pixel data that the client should put into its framebuffer.
type FramebufferUpdate struct {
	Rectangles []Rectangle
}

// Rectangle represents a rectangle of pixel data.
type Rectangle struct {
	X      uint16
	Y      uint16
	Width  uint16
	Height uint16
	Enc    Encoding
}

func (*FramebufferUpdate) Type() uint8 {
	return framebufferUpdateMsg
}

func (*FramebufferUpdate) Read(c *ClientConn, r io.Reader) (ServerMessage, error) {
	// Read off the padding
	var padding [1]byte
	if _, err := io.ReadFull(r, padding[:]); err != nil {
		return nil, err
	}

	var numRects uint16
	if err := binary.Read(r, binary.BigEndian, &numRects); err != nil {
		return nil, err
	}

	// Build the map of encodings supported
	encMap := make(map[int32]Encoding)
	for _, e := range c.Encodings {
		encMap[e.Type()] = e
	}

	// We must always support the raw encoding
	rawEnc := new(RawEncoding)
	encMap[rawEnc.Type()] = rawEnc

	rects := make([]Rectangle, numRects)
	for i := uint16(0); i < numRects; i++ {
		var encodingType int32

		rect := &rects[i]
		data := []interface{}{
			&rect.X,
			&rect.Y,
			&rect.Width,
			&rect.Height,
			&encodingType,
		}

		for _, val := range data {
			if err := binary.Read(r, binary.BigEndian, val); err != nil {
				return nil, err
			}
		}

		enc, ok := encMap[encodingType]
		if !ok {
			return nil, fmt.Errorf("unsupported encoding type: %d", encodingType)
		}

		var err error
		rect.Enc, err = enc.Read(c, rect, r)
		if err != nil {
			return nil, err
		}
	}

	return &FramebufferUpdate{rects}, nil
}

// SetColorMapEntries is sent by the server to set values into
// the color map. This message will automatically update the color map
// for the associated connection, but contains the color change data
// if the consumer wants to read it.
//
// See RFC 6143 Section 7.6.2

// Color represents a single color in a color map.
type Color struct {
	R, G, B uint16
}

type SetColorMapEntries struct {
	FirstColor uint16
	Colors     []Color
}

func (*SetColorMapEntries) Type() uint8 {
	return setColorMapEntriesMsg
}

func (*SetColorMapEntries) Read(c *ClientConn, r io.Reader) (ServerMessage, error) {
	// Read off the padding
	var padding [1]byte
	if _, err := io.ReadFull(r, padding[:]); err != nil {
		return nil, err
	}

	var result SetColorMapEntries
	if err := binary.Read(r, binary.BigEndian, &result.FirstColor); err != nil {
		return nil, err
	}

	var numColors uint16
	if err := binary.Read(r, binary.BigEndian, &numColors); err != nil {
		return nil, err
	}

	result.Colors = make([]Color, numColors)
	for i := uint16(0); i < numColors; i++ {

		color := &result.Colors[i]
		data := []interface{}{
			&color.R,
			&color.G,
			&color.B,
		}

		for _, val := range data {
			if err := binary.Read(r, binary.BigEndian, val); err != nil {
				return nil, err
			}
		}

		// Update the connection's color map
		c.ColorMap[result.FirstColor+i] = *color
	}

	return &result, nil
}

// Bell signals that an audible bell should be made on the client.
//
// See RFC 6143 Section 7.6.3
type Bell struct{}

func (*Bell) Type() uint8 {
	return bellMsg
}

func (*Bell) Read(*ClientConn, io.Reader) (ServerMessage, error) {
	return new(Bell), nil
}

// ServerCutText indicates the server has new text in the cut buffer.
//
// See RFC 6143 Section 7.6.4
type ServerCutText struct {
	Text string
}

func (*ServerCutText) Type() uint8 {
	return serverCutTextMsg
}

func (*ServerCutText) Read(c *ClientConn, r io.Reader) (ServerMessage, error) {
	// Read off the padding
	var padding [1]byte
	if _, err := io.ReadFull(r, padding[:]); err != nil {
		return nil, err
	}

	var textLength uint32
	if err := binary.Read(r, binary.BigEndian, &textLength); err != nil {
		return nil, err
	}

	textBytes := make([]uint8, textLength)
	if err := binary.Read(r, binary.BigEndian, &textBytes); err != nil {
		return nil, err
	}

	return &ServerCutText{string(textBytes)}, nil
}
<commit_msg>forced type of constants to uint8<commit_after>/*
Server implements RFC 6143 §7.6 Server-to-Client Messages.

See http://tools.ietf.org/html/rfc6143#section-7.6 for more info.
*/
package vnc

import (
	"encoding/binary"
	"fmt"
	"io"
)

const (
	framebufferUpdateMsg = uint8(iota)
	setColorMapEntriesMsg
	bellMsg
	serverCutTextMsg
)

// A ServerMessage implements a message sent from the server to the client.
type ServerMessage interface {
	// The type of the message that is sent down on the wire.
	Type() uint8

	// Read reads the contents of the message from the reader. At the point
	// this is called, the message type has already been read from the reader.
	// This should return a new ServerMessage that is the appropriate type.
	Read(*ClientConn, io.Reader) (ServerMessage, error)
}

// FramebufferUpdate consists of a sequence of rectangles of
// pixel data that the client should put into its framebuffer.
type FramebufferUpdate struct {
	Rectangles []Rectangle
}

// Rectangle represents a rectangle of pixel data.
type Rectangle struct {
	X      uint16
	Y      uint16
	Width  uint16
	Height uint16
	Enc    Encoding
}

func (*FramebufferUpdate) Type() uint8 {
	return framebufferUpdateMsg
}

func (*FramebufferUpdate) Read(c *ClientConn, r io.Reader) (ServerMessage, error) {
	// Read off the padding
	var padding [1]byte
	if _, err := io.ReadFull(r, padding[:]); err != nil {
		return nil, err
	}

	var numRects uint16
	if err := binary.Read(r, binary.BigEndian, &numRects); err != nil {
		return nil, err
	}

	// Build the map of encodings supported
	encMap := make(map[int32]Encoding)
	for _, e := range c.Encodings {
		encMap[e.Type()] = e
	}

	// We must always support the raw encoding
	rawEnc := new(RawEncoding)
	encMap[rawEnc.Type()] = rawEnc

	rects := make([]Rectangle, numRects)
	for i := uint16(0); i < numRects; i++ {
		var encodingType int32

		rect := &rects[i]
		data := []interface{}{
			&rect.X,
			&rect.Y,
			&rect.Width,
			&rect.Height,
			&encodingType,
		}

		for _, val := range data {
			if err := binary.Read(r, binary.BigEndian, val); err != nil {
				return nil, err
			}
		}

		enc, ok := encMap[encodingType]
		if !ok {
			return nil, fmt.Errorf("unsupported encoding type: %d", encodingType)
		}

		var err error
		rect.Enc, err = enc.Read(c, rect, r)
		if err != nil {
			return nil, err
		}
	}

	return &FramebufferUpdate{rects}, nil
}

// SetColorMapEntries is sent by the server to set values into
// the color map. This message will automatically update the color map
// for the associated connection, but contains the color change data
// if the consumer wants to read it.
//
// See RFC 6143 Section 7.6.2

// Color represents a single color in a color map.
type Color struct {
	R, G, B uint16
}

type SetColorMapEntries struct {
	FirstColor uint16
	Colors     []Color
}

func (*SetColorMapEntries) Type() uint8 {
	return setColorMapEntriesMsg
}

func (*SetColorMapEntries) Read(c *ClientConn, r io.Reader) (ServerMessage, error) {
	// Read off the padding
	var padding [1]byte
	if _, err := io.ReadFull(r, padding[:]); err != nil {
		return nil, err
	}

	var result SetColorMapEntries
	if err := binary.Read(r, binary.BigEndian, &result.FirstColor); err != nil {
		return nil, err
	}

	var numColors uint16
	if err := binary.Read(r, binary.BigEndian, &numColors); err != nil {
		return nil, err
	}

	result.Colors = make([]Color, numColors)
	for i := uint16(0); i < numColors; i++ {

		color := &result.Colors[i]
		data := []interface{}{
			&color.R,
			&color.G,
			&color.B,
		}

		for _, val := range data {
			if err := binary.Read(r, binary.BigEndian, val); err != nil {
				return nil, err
			}
		}

		// Update the connection's color map
		c.ColorMap[result.FirstColor+i] = *color
	}

	return &result, nil
}

// Bell signals that an audible bell should be made on the client.
//
// See RFC 6143 Section 7.6.3
type Bell struct{}

func (*Bell) Type() uint8 {
	return bellMsg
}

func (*Bell) Read(*ClientConn, io.Reader) (ServerMessage, error) {
	return new(Bell), nil
}

// ServerCutText indicates the server has new text in the cut buffer.
//
// See RFC 6143 Section 7.6.4
type ServerCutText struct {
	Text string
}

func (*ServerCutText) Type() uint8 {
	return serverCutTextMsg
}

func (*ServerCutText) Read(c *ClientConn, r io.Reader) (ServerMessage, error) {
	// Read off the padding
	var padding [1]byte
	if _, err := io.ReadFull(r, padding[:]); err != nil {
		return nil, err
	}

	var textLength uint32
	if err := binary.Read(r, binary.BigEndian, &textLength); err != nil {
		return nil, err
	}

	textBytes := make([]uint8, textLength)
	if err := binary.Read(r, binary.BigEndian, &textBytes); err != nil {
		return nil, err
	}

	return &ServerCutText{string(textBytes)}, nil
}
<|endoftext|>


<commit_before>// Server is the base tcpez server. It sets up a tcp listener on an address,
// and given a RequestHandler, it parses the tcpez protocol format and turns
// it into individual request/responses. Each Connection is handled on a
// seperate goroutine and pipelined requests are first parsed then farmed
// to seperate goroutines. Pipelined requests from the client are handled
// seamlessly this way, each seperate request is passed to its own RequestHandler
// with its own Span.
package tcpez

import (
	"bytes"
	"encoding/binary"
	"errors"
	"github.com/op/go-logging"
	"io"
	"net"
	"sync"
	"time"
)

var log = logging.MustGetLogger("tcpez")
var LogFormat = logging.MustStringFormatter("%{time:2006-01-02T15:04:05.999999999Z07:00} %{level} [%{module}] %{message}")

func init() {
	logging.SetLevel(logging.INFO, "tcpez")
	logging.SetFormatter(LogFormat)
}

// Server is the base struct that wraps the tcp listener and allows
// for setting the RequestHandler that takes each request and returns
// an response
type Server struct {
	// Address is the string address that the tcp listener is
	// bound to
	Address string
	// Handler is a value that responds to the RequestHandler
	// interface.
	Handler RequestHandler
	// StatsRecorder is the value that delivers stats to a collection
	// agent. This is the DebugStatsRecorder by default (does nothing)
	// but can be swapped to send to StatsD. This is passed to each
	// Span created and passed to the RequestHandler
	Stats StatsRecorder
	// UUIDGenerator generates UUIDs for each Span, by default this uses
	// a simple hash function but can be swapped out for something more
	// complex (a vector-clock style UUID generator for example)
	UUIDGenerator UUIDGenerator

	// The underlying TCP listener, access if you need to set timeouts, etc
	Conn *net.TCPListener
	// The number of open tcpez connections
	NumConnections int
	isClosed       bool
	closer         chan int
}

// RequestHandler is the basic interface for setting up the request handling
// logic of a tcpez server. The server handles all the request parsing and setup
// as well as the response encoding. All you have to do to create a working server
// is create an object that has a Respond() method that takes a byte slice (which
// is the request) and a Span (which allows you to track timings and meta data
// through the request) and then it returns a byte slice which is the response.
//
//        type MyHandler struct
//
//        func (h *MyHandler) Respond(req []byte, span *tcpez.Span) (res []byte, err error) {
//              if string(req) == "PING" {
//                  span.Attr("command", "PING")
//                  return []byte{"PONG"}, nil
//              }
//        }
//
type RequestHandler interface {
	Respond([]byte, *Span) ([]byte, error)
}

// NewServer is the tcpez server intializer. It only requires two parameters,
// an address to bind to (same format as net.ListenTCP) and a RequestHandler
// which serves the requests.
func NewServer(address string, handler RequestHandler) (s *Server, err error) {
	tcpAddr, err := net.ResolveTCPAddr("tcp", address)
	if err != nil {
		return nil, err
	}
	l, err := net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		return nil, err
	}

	return &Server{Address: address, Conn: l, Handler: handler, Stats: new(DebugStatsRecorder), UUIDGenerator: DefaultUUIDGenerator, closer: make(chan int)}, nil
}

// Start starts the Connection handling and request processing loop.
// This is a blocking operation and can be started in a goroutine.
func (s *Server) Start() {
	log.Debug("Listening on %s", s.Conn.Addr().String())
	for {
		if s.isClosed == true {
			break
		}
		clientConn, err := s.Conn.Accept()
		if err != nil {
			log.Warning(err.Error())
			break
		}
		s.NumConnections++
		go s.clientCloser(clientConn)
		go s.handle(clientConn)
	}
	log.Debug("Closing %s", s.Conn.Addr().String())
}

// Close closes the server listener to any more Connections
func (s *Server) Close() (err error) {
	if s.isClosed == false {
		err = s.Conn.Close()
		s.isClosed = true
		for i := 0; i < s.NumConnections; i++ {
			s.closer <- 1
		}
		return
	}
	return errors.New("Closing already closed Connection")
}

func (s *Server) handle(clientConn net.Conn) {
	log.Debug("[tcpez] New client(%s)", clientConn.RemoteAddr())
	for {
		// Timeout the connection after 5 mins
		clientConn.SetReadDeadline(time.Now().Add(5 * time.Minute))
		header, response, err := s.readHeaderAndHandleRequest(clientConn)
		if err != nil {
			if closableError(err) {
				// EOF the client has disconnected
				break
			}
			log.Error(err.Error())
			s.Stats.Increment("operation.failure")
			return
		}
		err = s.sendResponse(clientConn, header, response)
		if err != nil {
			if closableError(err) {
				// EOF the client has disconnected
				break
			}
			log.Error(err.Error())
			s.Stats.Increment("operation.failure")
			return
		}
		s.Stats.Increment("operation.success")
	}
	clientConn.Close()
	s.NumConnections--
}

func closableError(err error) bool {
	return err == io.EOF || err == io.ErrClosedPipe || err.(net.Error).Timeout() == true
}

// clientCloser closes client connections when the server is closed
func (s *Server) clientCloser(clientConn net.Conn) {
	<-s.closer
	log.Debug("Closing the client connection %v", clientConn)
	clientConn.Close()
}

func (s *Server) readHeaderAndHandleRequest(buf io.Reader) (header int32, response []byte, err error) {
	var size int32
	err = binary.Read(buf, binary.BigEndian, &size)
	if err != nil {
		return 0, nil, err
	}
	if size < 0 {
		// this is a pipelined request
		var wg sync.WaitGroup
		count := -size
		requests := make([][]byte, count)
		responses := make([][]byte, count)
		for r := 0; int32(r) < count; r++ {
			request, err := s.parseRequest(buf, 0)
			if err == nil {
				requests[r] = request
				wg.Add(1)
				go func(index int) {
					res, err := s.handleRequest(requests[index], true)
					if err == nil {
						responses[index] = res
					}
					wg.Done()
				}(r)
			}
		}
		wg.Wait()
		output := bytes.NewBuffer(nil)
		for j := 0; int32(j) < count; j++ {
			length := int32(len(responses[j]))
			err = binary.Write(output, binary.BigEndian, length)
			if err == nil {
				output.Write(responses[j])
			}
		}
		return int32(-count), output.Bytes(), err
	} else {
		request, err := s.parseRequest(buf, size)
		if err != nil {
			return 0, nil, err
		}
		response, err := s.handleRequest(request, false)
		if err != nil {
			return 0, nil, err
		}
		return int32(len(response)), response, nil
	}
	return
}

func (s *Server) sendResponse(w io.Writer, header int32, data []byte) (err error) {
	err = binary.Write(w, binary.BigEndian, header)
	if err != nil {
		return err
	}
	_, err = w.Write(data)
	return
}

func (s *Server) parseRequest(buf io.Reader, size int32) (request []byte, err error) {
	if size == int32(0) {
		err = binary.Read(buf, binary.BigEndian, &size)
		if err != nil {
			return nil, err
		}
	}
	request = make([]byte, size)
	_, err = io.ReadFull(buf, request)
	if err != nil {
		return nil, err
	}
	return
}

func (s *Server) handleRequest(request []byte, multi bool) (response []byte, err error) {
	span := NewSpan(s.UUIDGenerator())
	if multi == true {
		span.Attr("multi", "true")
	}
	span.Stats = s.Stats
	span.Start("duration")
	span.Add("num_connections", int64(s.NumConnections))
	response, err = s.Handler.Respond(request, span)
	span.Finish("duration")
	log.Info("%s", span.JSON())
	span.Record()
	return
}
<commit_msg>Reimplement clientCloser so that it doesnt create a goroutine per conn<commit_after>// Server is the base tcpez server. It sets up a tcp listener on an address,
// and given a RequestHandler, it parses the tcpez protocol format and turns
// it into individual request/responses. Each Connection is handled on a
// seperate goroutine and pipelined requests are first parsed then farmed
// to seperate goroutines. Pipelined requests from the client are handled
// seamlessly this way, each seperate request is passed to its own RequestHandler
// with its own Span.
package tcpez

import (
	"bytes"
	"encoding/binary"
	"errors"
	"github.com/op/go-logging"
	"io"
	"net"
	"sync"
	"time"
)

var log = logging.MustGetLogger("tcpez")
var LogFormat = logging.MustStringFormatter("%{time:2006-01-02T15:04:05.999999999Z07:00} %{level} [%{module}] %{message}")

func init() {
	logging.SetLevel(logging.INFO, "tcpez")
	logging.SetFormatter(LogFormat)
}

// Server is the base struct that wraps the tcp listener and allows
// for setting the RequestHandler that takes each request and returns
// an response
type Server struct {
	// Address is the string address that the tcp listener is
	// bound to
	Address string
	// Handler is a value that responds to the RequestHandler
	// interface.
	Handler RequestHandler
	// StatsRecorder is the value that delivers stats to a collection
	// agent. This is the DebugStatsRecorder by default (does nothing)
	// but can be swapped to send to StatsD. This is passed to each
	// Span created and passed to the RequestHandler
	Stats StatsRecorder
	// UUIDGenerator generates UUIDs for each Span, by default this uses
	// a simple hash function but can be swapped out for something more
	// complex (a vector-clock style UUID generator for example)
	UUIDGenerator UUIDGenerator

	// The underlying TCP listener, access if you need to set timeouts, etc
	Conn *net.TCPListener

	isClosed    bool
	connId      int
	clientConns map[int]net.Conn
}

// RequestHandler is the basic interface for setting up the request handling
// logic of a tcpez server. The server handles all the request parsing and setup
// as well as the response encoding. All you have to do to create a working server
// is create an object that has a Respond() method that takes a byte slice (which
// is the request) and a Span (which allows you to track timings and meta data
// through the request) and then it returns a byte slice which is the response.
//
//        type MyHandler struct
//
//        func (h *MyHandler) Respond(req []byte, span *tcpez.Span) (res []byte, err error) {
//              if string(req) == "PING" {
//                  span.Attr("command", "PING")
//                  return []byte{"PONG"}, nil
//              }
//        }
//
type RequestHandler interface {
	Respond([]byte, *Span) ([]byte, error)
}

// NewServer is the tcpez server intializer. It only requires two parameters,
// an address to bind to (same format as net.ListenTCP) and a RequestHandler
// which serves the requests.
func NewServer(address string, handler RequestHandler) (s *Server, err error) {
	tcpAddr, err := net.ResolveTCPAddr("tcp", address)
	if err != nil {
		return nil, err
	}
	l, err := net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		return nil, err
	}

	return &Server{Address: address, Conn: l, Handler: handler, Stats: new(DebugStatsRecorder), UUIDGenerator: DefaultUUIDGenerator, clientConns: make(map[int]net.Conn)}, nil
}

// Start starts the Connection handling and request processing loop.
// This is a blocking operation and can be started in a goroutine.
func (s *Server) Start() {
	log.Debug("Listening on %s", s.Conn.Addr().String())
	for {
		if s.isClosed == true {
			break
		}
		clientConn, err := s.Conn.Accept()
		if err != nil {
			log.Warning(err.Error())
			break
		}
		s.connId++
		go s.handle(clientConn, s.connId)
	}
	log.Debug("Closing %s", s.Conn.Addr().String())
}

func (s *Server) NumConnections() int {
	return len(s.clientConns)
}

// Close closes the server listener to any more Connections
func (s *Server) Close() (err error) {
	if s.isClosed == false {
		err = s.Conn.Close()
		s.isClosed = true
		for id, conn := range s.clientConns {
			delete(s.clientConns, id)
			conn.Close()
		}
		return
	}
	return errors.New("Closing already closed Connection")
}

func (s *Server) handle(clientConn net.Conn, id int) {
	log.Debug("[tcpez] New client(%s)", clientConn.RemoteAddr())
	s.clientConns[id] = clientConn
	for {
		// Timeout the connection after 5 mins
		clientConn.SetReadDeadline(time.Now().Add(5 * time.Minute))
		header, response, err := s.readHeaderAndHandleRequest(clientConn)
		if err != nil {
			if closableError(err) {
				// EOF the client has disconnected
				break
			}
			log.Error(err.Error())
			s.Stats.Increment("operation.failure")
			return
		}
		err = s.sendResponse(clientConn, header, response)
		if err != nil {
			if closableError(err) {
				// EOF the client has disconnected
				break
			}
			log.Error(err.Error())
			s.Stats.Increment("operation.failure")
			return
		}
		s.Stats.Increment("operation.success")
	}
	clientConn.Close()
	delete(s.clientConns, id)
}

func closableError(err error) bool {
	return err == io.EOF || err == io.ErrClosedPipe || err.(net.Error).Timeout() == true
}

func (s *Server) readHeaderAndHandleRequest(buf io.Reader) (header int32, response []byte, err error) {
	var size int32
	err = binary.Read(buf, binary.BigEndian, &size)
	if err != nil {
		return 0, nil, err
	}
	if size < 0 {
		// this is a pipelined request
		var wg sync.WaitGroup
		count := -size
		requests := make([][]byte, count)
		responses := make([][]byte, count)
		for r := 0; int32(r) < count; r++ {
			request, err := s.parseRequest(buf, 0)
			if err == nil {
				requests[r] = request
				wg.Add(1)
				go func(index int) {
					res, err := s.handleRequest(requests[index], true)
					if err == nil {
						responses[index] = res
					}
					wg.Done()
				}(r)
			}
		}
		wg.Wait()
		output := bytes.NewBuffer(nil)
		for j := 0; int32(j) < count; j++ {
			length := int32(len(responses[j]))
			err = binary.Write(output, binary.BigEndian, length)
			if err == nil {
				output.Write(responses[j])
			}
		}
		return int32(-count), output.Bytes(), err
	} else {
		request, err := s.parseRequest(buf, size)
		if err != nil {
			return 0, nil, err
		}
		response, err := s.handleRequest(request, false)
		if err != nil {
			return 0, nil, err
		}
		return int32(len(response)), response, nil
	}
	return
}

func (s *Server) sendResponse(w io.Writer, header int32, data []byte) (err error) {
	err = binary.Write(w, binary.BigEndian, header)
	if err != nil {
		return err
	}
	_, err = w.Write(data)
	return
}

func (s *Server) parseRequest(buf io.Reader, size int32) (request []byte, err error) {
	if size == int32(0) {
		err = binary.Read(buf, binary.BigEndian, &size)
		if err != nil {
			return nil, err
		}
	}
	request = make([]byte, size)
	_, err = io.ReadFull(buf, request)
	if err != nil {
		return nil, err
	}
	return
}

func (s *Server) handleRequest(request []byte, multi bool) (response []byte, err error) {
	span := NewSpan(s.UUIDGenerator())
	if multi == true {
		span.Attr("multi", "true")
	}
	span.Stats = s.Stats
	span.Start("duration")
	span.Add("num_connections", int64(s.NumConnections()))
	response, err = s.Handler.Respond(request, span)
	span.Finish("duration")
	log.Info("%s", span.JSON())
	span.Record()
	return
}
<|endoftext|>


<commit_before>// Copyright (C) 2016 Kale Blankenship. All rights reserved.
// This software may be modified and distributed under the terms
// of the MIT license.  See the LICENSE file for details

package trivialt

import "net"

// Server contains the configuration to run a TFTP server.
//
// A ReadHandler, WriteHandler, or both can be registered to the server. If one
// of the handlers isn't registered, the server will return errors to clients
// attempting to use them.
type Server struct {
	log     *logger
	net     string
	addrStr string
	addr    *net.UDPAddr
	conn    *net.UDPConn
	close   bool

	retransmit int // Per-packet retransmission limit

	rh ReadHandler
	wh WriteHandler
}

// NewServer returns a configured Server.
//
// Addr is the network address to listen on and is in the form "host:port".
// If a no host is given the server will listen on all interfaces.
//
// Any number of ServerOpts can be provided to configure optional values.
func NewServer(addr string, opts ...ServerOpt) (*Server, error) {
	s := &Server{
		log:        newLogger("server"),
		net:        defaultUDPNet,
		addrStr:    addr,
		retransmit: defaultRetransmit,
	}

	for _, opt := range opts {
		if err := opt(s); err != nil {
			return nil, err
		}
	}

	return s, nil
}

// Addr is the network address of the server. It is available
// after the server has been started.
func (s *Server) Addr() (*net.UDPAddr, error) {
	if s.conn == nil {
		return nil, ErrAddressNotAvailable
	}
	return s.conn.LocalAddr().(*net.UDPAddr), nil
}

// ReadHandler registers a ReadHandler for the server.
func (s *Server) ReadHandler(rh ReadHandler) {
	s.rh = rh
}

// WriteHandler registers a WriteHandler for the server.
func (s *Server) WriteHandler(wh WriteHandler) {
	s.wh = wh
}

// Serve starts the server using an existing UDPConn.
func (s *Server) Serve(conn *net.UDPConn) error {
	if s.rh == nil && s.wh == nil {
		return ErrNoRegisteredHandlers
	}

	s.conn = conn

	buf := make([]byte, 1024)
	for {
		numBytes, addr, err := conn.ReadFromUDP(buf)
		if err != nil {
			if s.close {
				return nil
			}
			return wrapError(err, "reading from conn")
		}
		bufCopy := make([]byte, numBytes)
		copy(bufCopy, buf)
		go s.dispatchRequest(addr, bufCopy)
	}
}

// Close stops the server and closes the network connection.
func (s *Server) Close() error {
	s.close = true
	return s.conn.Close()
}

// dispatchRequest parses incoming requests and executes the corresponding
// handler, if one is registered. If a handler is not registered the server
// sends an error to the client.
func (s *Server) dispatchRequest(addr *net.UDPAddr, b []byte) {
	conn, err := newConn(s.net, "", addr) // Use empty mode until request has been parsed.
	if err != nil {
		s.log.err("Received error opening connection for new request: %v", err)
		return
	}
	defer errorDefer(conn.Close, s.log, "error closing network connection in dispath")

	// Set retransmit
	conn.retransmit = s.retransmit

	// Validate request datagram
	conn.rx.setBytes(b)
	if err := conn.rx.validate(); err != nil {
		s.log.debug("Error decoding new request: %v", err)
		return
	}
	s.log.debug("New request from %v: %s", addr, conn.rx)

	// Set mode from request
	conn.mode = conn.rx.mode()

	switch conn.rx.opcode() {
	case opCodeRRQ:
		// Check for handler
		if s.rh == nil {
			s.log.debug("No read handler registered.")
			conn.sendError(ErrCodeIllegalOperation, "Server does not support read requests.")
			return
		}

		// Create request
		w := &readRequest{conn: conn, name: conn.rx.filename()}

		// execute handler
		s.rh.ServeTFTP(w)
	case opCodeWRQ:
		// Check for handler
		if s.wh == nil {
			s.log.debug("No write handler registered.")
			conn.sendError(ErrCodeIllegalOperation, "Server does not support write requests.")
			return
		}

		// Create request
		w := &writeRequest{conn: conn, name: conn.rx.filename()}

		// parse options to get size
		conn.log.trace("performing write setup")
		if err := conn.readSetup(); err != nil {
			conn.err = err
		}

		s.wh.ReceiveTFTP(w)
	default:
		s.log.debug("Unexpected Request")
	}
}

// ListenAndServe starts a configured server.
func (s *Server) ListenAndServe() error {
	addr, err := net.ResolveUDPAddr(s.net, s.addrStr)
	if err != nil {
		return wrapError(err, "resolving server address")
	}
	s.addr = addr

	conn, err := net.ListenUDP(s.net, s.addr)
	if err != nil {
		return wrapError(err, "opening network connection")
	}

	return wrapError(s.Serve(conn), "serving tftp")
}

// ListenAndServe creates and starts a Server with default options.
func ListenAndServe(addr string, rh ReadHandler, wh WriteHandler) error {
	s, err := NewServer(addr)
	if err != nil {
		return wrapError(err, "creaing new server")
	}

	s.rh = rh
	s.wh = wh

	return wrapError(s.ListenAndServe(), "starting server")
}

// ServerOpt is a function that configures a Server.
type ServerOpt func(*Server) error

// ServerNet configures the network a server listens on.
// Must be one of: udp, udp4, udp6.
//
// Default: udp.
func ServerNet(net string) ServerOpt {
	return func(s *Server) error {
		if net != "udp" && net != "udp4" && net != "udp6" {
			return ErrInvalidNetwork
		}
		s.net = net
		return nil
	}
}

// ServerRetransmit configures the per-packet retransmission limit for all requests.
//
// Default: 10.
func ServerRetransmit(i int) ServerOpt {
	return func(s *Server) error {
		if i < 0 {
			return ErrInvalidRetransmit
		}
		s.retransmit = i
		return nil
	}
}
<commit_msg>Removed unused ListenAndServer func<commit_after>// Copyright (C) 2016 Kale Blankenship. All rights reserved.
// This software may be modified and distributed under the terms
// of the MIT license.  See the LICENSE file for details

package trivialt

import "net"

// Server contains the configuration to run a TFTP server.
//
// A ReadHandler, WriteHandler, or both can be registered to the server. If one
// of the handlers isn't registered, the server will return errors to clients
// attempting to use them.
type Server struct {
	log     *logger
	net     string
	addrStr string
	addr    *net.UDPAddr
	conn    *net.UDPConn
	close   bool

	retransmit int // Per-packet retransmission limit

	rh ReadHandler
	wh WriteHandler
}

// NewServer returns a configured Server.
//
// Addr is the network address to listen on and is in the form "host:port".
// If a no host is given the server will listen on all interfaces.
//
// Any number of ServerOpts can be provided to configure optional values.
func NewServer(addr string, opts ...ServerOpt) (*Server, error) {
	s := &Server{
		log:        newLogger("server"),
		net:        defaultUDPNet,
		addrStr:    addr,
		retransmit: defaultRetransmit,
	}

	for _, opt := range opts {
		if err := opt(s); err != nil {
			return nil, err
		}
	}

	return s, nil
}

// Addr is the network address of the server. It is available
// after the server has been started.
func (s *Server) Addr() (*net.UDPAddr, error) {
	if s.conn == nil {
		return nil, ErrAddressNotAvailable
	}
	return s.conn.LocalAddr().(*net.UDPAddr), nil
}

// ReadHandler registers a ReadHandler for the server.
func (s *Server) ReadHandler(rh ReadHandler) {
	s.rh = rh
}

// WriteHandler registers a WriteHandler for the server.
func (s *Server) WriteHandler(wh WriteHandler) {
	s.wh = wh
}

// Serve starts the server using an existing UDPConn.
func (s *Server) Serve(conn *net.UDPConn) error {
	if s.rh == nil && s.wh == nil {
		return ErrNoRegisteredHandlers
	}

	s.conn = conn

	buf := make([]byte, 1024)
	for {
		numBytes, addr, err := conn.ReadFromUDP(buf)
		if err != nil {
			if s.close {
				return nil
			}
			return wrapError(err, "reading from conn")
		}
		bufCopy := make([]byte, numBytes)
		copy(bufCopy, buf)
		go s.dispatchRequest(addr, bufCopy)
	}
}

// Close stops the server and closes the network connection.
func (s *Server) Close() error {
	s.close = true
	return s.conn.Close()
}

// dispatchRequest parses incoming requests and executes the corresponding
// handler, if one is registered. If a handler is not registered the server
// sends an error to the client.
func (s *Server) dispatchRequest(addr *net.UDPAddr, b []byte) {
	conn, err := newConn(s.net, "", addr) // Use empty mode until request has been parsed.
	if err != nil {
		s.log.err("Received error opening connection for new request: %v", err)
		return
	}
	defer errorDefer(conn.Close, s.log, "error closing network connection in dispath")

	// Set retransmit
	conn.retransmit = s.retransmit

	// Validate request datagram
	conn.rx.setBytes(b)
	if err := conn.rx.validate(); err != nil {
		s.log.debug("Error decoding new request: %v", err)
		return
	}
	s.log.debug("New request from %v: %s", addr, conn.rx)

	// Set mode from request
	conn.mode = conn.rx.mode()

	switch conn.rx.opcode() {
	case opCodeRRQ:
		// Check for handler
		if s.rh == nil {
			s.log.debug("No read handler registered.")
			conn.sendError(ErrCodeIllegalOperation, "Server does not support read requests.")
			return
		}

		// Create request
		w := &readRequest{conn: conn, name: conn.rx.filename()}

		// execute handler
		s.rh.ServeTFTP(w)
	case opCodeWRQ:
		// Check for handler
		if s.wh == nil {
			s.log.debug("No write handler registered.")
			conn.sendError(ErrCodeIllegalOperation, "Server does not support write requests.")
			return
		}

		// Create request
		w := &writeRequest{conn: conn, name: conn.rx.filename()}

		// parse options to get size
		conn.log.trace("performing write setup")
		if err := conn.readSetup(); err != nil {
			conn.err = err
		}

		s.wh.ReceiveTFTP(w)
	default:
		s.log.debug("Unexpected Request")
	}
}

// ListenAndServe starts a configured server.
func (s *Server) ListenAndServe() error {
	addr, err := net.ResolveUDPAddr(s.net, s.addrStr)
	if err != nil {
		return wrapError(err, "resolving server address")
	}
	s.addr = addr

	conn, err := net.ListenUDP(s.net, s.addr)
	if err != nil {
		return wrapError(err, "opening network connection")
	}

	return wrapError(s.Serve(conn), "serving tftp")
}

// ServerOpt is a function that configures a Server.
type ServerOpt func(*Server) error

// ServerNet configures the network a server listens on.
// Must be one of: udp, udp4, udp6.
//
// Default: udp.
func ServerNet(net string) ServerOpt {
	return func(s *Server) error {
		if net != "udp" && net != "udp4" && net != "udp6" {
			return ErrInvalidNetwork
		}
		s.net = net
		return nil
	}
}

// ServerRetransmit configures the per-packet retransmission limit for all requests.
//
// Default: 10.
func ServerRetransmit(i int) ServerOpt {
	return func(s *Server) error {
		if i < 0 {
			return ErrInvalidRetransmit
		}
		s.retransmit = i
		return nil
	}
}
<|endoftext|>


<commit_before>// Copyright 2020 The Prometheus 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.

// +build !nonetdev

package collector

import (
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"strconv"

	"golang.org/x/sys/unix"
	"regexp"
	"unsafe"
)

func getNetDevStats(ignore *regexp.Regexp, accept *regexp.Regexp, logger log.Logger) (map[string]map[string]string, error) {
	netDev := map[string]map[string]string{}

	mib := [6]_C_int{unix.CTL_NET, unix.AF_ROUTE, 0, 0, unix.NET_RT_IFLIST, 0}
	buf, err := sysctl(mib[:])
	if err != nil {
		return nil, err
	}
	n := uintptr(len(buf))
	index := uintptr(unsafe.Pointer(&buf[0]))
	next := uintptr(0)

	var rtm *unix.RtMsghdr

	for next = index; next < (index + n); next += uintptr(rtm.Msglen) {
		rtm = (*unix.RtMsghdr)(unsafe.Pointer(next))
		if rtm.Version != unix.RTM_VERSION || rtm.Type != unix.RTM_IFINFO {
			continue
		}
		ifm := (*unix.IfMsghdr)(unsafe.Pointer(next))
		if ifm.Addrs&unix.RTA_IFP == 0 {
			continue
		}
		dl := (*unix.RawSockaddrDatalink)(unsafe.Pointer(next + uintptr(rtm.Hdrlen)))
		if dl.Family != unix.AF_LINK {
			continue
		}
		data := ifm.Data
		dev := int8ToString(dl.Data[:dl.Nlen])
		if ignore != nil && ignore.MatchString(dev) {
			level.Debug(logger).Log("msg", "Ignoring device", "device", dev)
			continue
		}
		if accept != nil && !accept.MatchString(dev) {
			level.Debug(logger).Log("msg", "Ignoring device", "device", dev)
			continue
		}

		netDev[dev] = map[string]string{
			"receive_packets":    strconv.Itoa(int(data.Ipackets)),
			"transmit_packets":   strconv.Itoa(int(data.Opackets)),
			"receive_errs":       strconv.Itoa(int(data.Ierrors)),
			"transmit_errs":      strconv.Itoa(int(data.Oerrors)),
			"receive_bytes":      strconv.Itoa(int(data.Ibytes)),
			"transmit_bytes":     strconv.Itoa(int(data.Obytes)),
			"receive_multicast":  strconv.Itoa(int(data.Imcasts)),
			"transmit_multicast": strconv.Itoa(int(data.Omcasts)),
			"receive_drop":       strconv.Itoa(int(data.Iqdrops)),
		}
	}
	return netDev, nil
}
<commit_msg>fixed build<commit_after>// Copyright 2020 The Prometheus 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.

// +build !nonetdev

package collector

import (
	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"

	"golang.org/x/sys/unix"
	"regexp"
	"unsafe"
)

func getNetDevStats(ignore *regexp.Regexp, accept *regexp.Regexp, logger log.Logger) (netDevStats, error) {
	netDev := netDevStats{}

	mib := [6]_C_int{unix.CTL_NET, unix.AF_ROUTE, 0, 0, unix.NET_RT_IFLIST, 0}
	buf, err := sysctl(mib[:])
	if err != nil {
		return nil, err
	}
	n := uintptr(len(buf))
	index := uintptr(unsafe.Pointer(&buf[0]))
	next := uintptr(0)

	var rtm *unix.RtMsghdr

	for next = index; next < (index + n); next += uintptr(rtm.Msglen) {
		rtm = (*unix.RtMsghdr)(unsafe.Pointer(next))
		if rtm.Version != unix.RTM_VERSION || rtm.Type != unix.RTM_IFINFO {
			continue
		}
		ifm := (*unix.IfMsghdr)(unsafe.Pointer(next))
		if ifm.Addrs&unix.RTA_IFP == 0 {
			continue
		}
		dl := (*unix.RawSockaddrDatalink)(unsafe.Pointer(next + uintptr(rtm.Hdrlen)))
		if dl.Family != unix.AF_LINK {
			continue
		}
		data := ifm.Data
		dev := int8ToString(dl.Data[:dl.Nlen])
		if ignore != nil && ignore.MatchString(dev) {
			level.Debug(logger).Log("msg", "Ignoring device", "device", dev)
			continue
		}
		if accept != nil && !accept.MatchString(dev) {
			level.Debug(logger).Log("msg", "Ignoring device", "device", dev)
			continue
		}

		netDev[dev] = map[string]uint64{
			"receive_packets":    data.Ipackets,
			"transmit_packets":   data.Opackets,
			"receive_errs":       data.Ierrors,
			"transmit_errs":      data.Oerrors,
			"receive_bytes":      data.Ibytes,
			"transmit_bytes":     data.Obytes,
			"receive_multicast":  data.Imcasts,
			"transmit_multicast": data.Omcasts,
			"receive_drop":       data.Iqdrops,
		}
	}
	return netDev, nil
}
<|endoftext|>


<commit_before>package incept

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

	"github.com/synapse-garden/sg-proto/store"
	"github.com/synapse-garden/sg-proto/users"

	"github.com/boltdb/bolt"
	uuid "github.com/satori/go.uuid"
)

var (
	TicketBucket = store.Bucket("tickets")

	masterUUID uuid.UUID
)

func init() {
	masterUUID = uuid.NewV4()
}

type Ticket uuid.UUID

func (t Ticket) Bytes() []byte  { return uuid.UUID(t).Bytes() }
func (t Ticket) String() string { return uuid.UUID(t).String() }

func NewTicket(tx *bolt.Tx) (Ticket, error) {
	u := uuid.NewV5(masterUUID, time.Now().String())
	if err := store.Put(TicketBucket, u.Bytes(), nil)(tx); err != nil {
		return Ticket(uuid.Nil), err
	}
	return Ticket(u), nil
}

func CheckTicket(key Ticket) func(*bolt.Tx) error {
	return func(tx *bolt.Tx) error {
		bs := tx.Bucket(TicketBucket).Get(key.Bytes())
		if bs == nil {
			return store.MissingError(key.Bytes())
		}
		return nil
	}
}

func DeleteKey(key Ticket) func(*bolt.Tx) error {
	return func(tx *bolt.Tx) error {
		return tx.Bucket(TicketBucket).Delete(key.Bytes())
	}
}

// Incept checks that the given Ticket exists, and that the given User
// does not (by name.)  Then it tries to create the given user,
// and delete the given key.  Any error will cause this to roll back.
func Incept(
	w io.Writer,
	key Ticket,
	u *users.User,
	db *bolt.DB,
) error {
	kbs, name := key.Bytes(), []byte(u.Name)
	// Check if given ticket exists (nil => it exists)
	if err := db.View(store.Wrap(
		store.CheckExists(TicketBucket, kbs),
		store.CheckNotExist(users.UserBucket, name),
	)); err != nil {
		return err
	}
	// Create user or fail
	if err := db.Update(store.Wrap(
		store.Delete(TicketBucket, kbs),
		users.Create(u),
	)); err != nil {
		return err
	}

	return json.NewEncoder(w).Encode(u)
}
<commit_msg>Add TODO to incept Windows bug code<commit_after>package incept

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

	"github.com/synapse-garden/sg-proto/store"
	"github.com/synapse-garden/sg-proto/users"

	"github.com/boltdb/bolt"
	uuid "github.com/satori/go.uuid"
)

var (
	TicketBucket = store.Bucket("tickets")

	masterUUID uuid.UUID
)

func init() {
	masterUUID = uuid.NewV4()
}

type Ticket uuid.UUID

func (t Ticket) Bytes() []byte  { return uuid.UUID(t).Bytes() }
func (t Ticket) String() string { return uuid.UUID(t).String() }

func NewTicket(tx *bolt.Tx) (Ticket, error) {
	// http://stackoverflow.com/a/14611355
	// This is completely unacceptable and must be resolved.  Two
	// Tickets MUST ***NEVER*** BE THE SAME and this should never
	// have been considered as an option.
	u := uuid.NewV5(masterUUID, time.Now().String())
	if err := store.Put(TicketBucket, u.Bytes(), nil)(tx); err != nil {
		return Ticket(uuid.Nil), err
	}
	return Ticket(u), nil
}

func CheckTicket(key Ticket) func(*bolt.Tx) error {
	return func(tx *bolt.Tx) error {
		bs := tx.Bucket(TicketBucket).Get(key.Bytes())
		if bs == nil {
			return store.MissingError(key.Bytes())
		}
		return nil
	}
}

func DeleteKey(key Ticket) func(*bolt.Tx) error {
	return func(tx *bolt.Tx) error {
		return tx.Bucket(TicketBucket).Delete(key.Bytes())
	}
}

// Incept checks that the given Ticket exists, and that the given User
// does not (by name.)  Then it tries to create the given user,
// and delete the given key.  Any error will cause this to roll back.
func Incept(
	w io.Writer,
	key Ticket,
	u *users.User,
	db *bolt.DB,
) error {
	kbs, name := key.Bytes(), []byte(u.Name)
	// Check if given ticket exists (nil => it exists)
	if err := db.View(store.Wrap(
		store.CheckExists(TicketBucket, kbs),
		store.CheckNotExist(users.UserBucket, name),
	)); err != nil {
		return err
	}
	// Create user or fail
	if err := db.Update(store.Wrap(
		store.Delete(TicketBucket, kbs),
		users.Create(u),
	)); err != nil {
		return err
	}

	return json.NewEncoder(w).Encode(u)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"text/template"

	"github.com/golang/protobuf/proto"
	descriptor "google/protobuf"
	plugin "google/protobuf/compiler"

	"pixur.org/pixur/schema/db/model"
)

const (
	tpl = `
package tables

import (
  "database/sql"

  "github.com/golang/protobuf/proto"

  "pixur.org/pixur/schema/db"
)

type Job struct {
  Tx *sql.Tx
}

var SqlTables = []string{
{{range .}}
  {{- .SqlString -}}
{{end -}}
}

{{range .}}
{{.ScanString}}
{{end -}}
`
)

type columnType string

var (
	int16Column columnType = "smallint"
	int32Column columnType = "integer"
	int64Column columnType = "bigint"
	bytesColumn columnType = "bytea"
)

type keyType string

var (
	primaryKey keyType = "PRIMARY KEY"
	uniqueKey  keyType = "UNIQUE"
	indexKey   keyType = "INDEX"
)

type column struct {
	Name       string
	ColumnType columnType
	field      *descriptor.FieldDescriptorProto
}

type index struct {
	Name    string
	KeyType keyType
	Columns []string
}

type table struct {
	Name    string
	GoType  string
	Columns []column
	Indexes []index
	msg     *descriptor.DescriptorProto
}

func (t table) SqlString() string {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, `  "CREATE TABLE \"%s\" (" +`, t.Name)
	buf.WriteRune('\n')
	for _, col := range t.Columns {
		fmt.Fprintf(&buf, `    "\"%s\" %s NOT NULL, " +`, col.Name, col.ColumnType)
		buf.WriteRune('\n')
	}
	var inlineIndexes []index
	var indexes []index
	for _, idx := range t.Indexes {
		if idx.KeyType == indexKey {
			indexes = append(indexes, idx)
		} else {
			inlineIndexes = append(inlineIndexes, idx)
		}
	}

	for i, idx := range inlineIndexes {
		switch idx.KeyType {
		case indexKey:
			continue
		}
		var cols []string
		for _, col := range idx.Columns {
			cols = append(cols, fmt.Sprintf(`\"%s\"`, col))
		}
		last := ", "
		if i == len(inlineIndexes)-1 {
			last = ""
		}
		fmt.Fprintf(&buf, `    "%s(%s)%s" +`, idx.KeyType, strings.Join(cols, ", "), last)
		buf.WriteRune('\n')
	}
	buf.WriteString("  \");\",\n")
	for _, idx := range indexes {
		var cols []string
		for _, col := range idx.Columns {
			cols = append(cols, fmt.Sprintf(`\"%s\"`, col))
		}
		fmt.Fprintf(&buf, `  "CREATE INDEX \"%s\" ON \"%s\" (%s);",`,
			idx.Name, t.Name, strings.Join(cols, ", "))
		buf.WriteRune('\n')
	}
	return buf.String()
}

func (t table) FindString() string {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, `func (j Job) Find%s(opts Opts) ([]%s, error) {`, t.Name, t.GoType)
	buf.WriteRune('\n')
	fmt.Fprintf(&buf, "\t"+`return db.Scan(j.Tx, "%s", opts, func(data []byte) error {`, t.Name)
	buf.WriteRune('\n')
	fmt.Fprintf(&buf, "\t\t"+`var pb %s`, t.GoType)
	buf.WriteRune('\n')
	buf.WriteString("\t\tif err := proto.Unmarshal(data, &pb); err != nil {\n")
	buf.WriteString("\t\t\treturn err\n")
	buf.WriteString("\t\t}\n")
	buf.WriteString("\treturn cb(pb)\n")
	buf.WriteRune('}')
	return buf.String()
}

func (t table) ScanString() string {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, `func (j Job) Scan%s(opts Opts, cb func(%s) error) error {`, t.Name, t.GoType)
	buf.WriteRune('\n')
	fmt.Fprintf(&buf, "\t"+`return db.Scan(j.Tx, "%s", opts, func(data []byte) error {`, t.Name)
	buf.WriteRune('\n')
	fmt.Fprintf(&buf, "\t\t"+`var pb %s`, t.GoType)
	buf.WriteRune('\n')
	buf.WriteString("\t\tif err := proto.Unmarshal(data, &pb); err != nil {\n")
	buf.WriteString("\t\t\treturn err\n")
	buf.WriteString("\t\t}\n")
	buf.WriteString("\treturn cb(pb)\n")
	buf.WriteRune('}')
	return buf.String()
}

func run() error {
	fds, err := getRequest()
	if err != nil {
		return err
	}

	log.Println(fds.String())

	var tables []table
	for _, fd := range fds.ProtoFile {
		for _, msg := range fd.MessageType {
			if msg.Options == nil || !proto.HasExtension(msg.Options, model.E_TabOpts) {
				continue
			}
			opts, err := proto.GetExtension(msg.Options, model.E_TabOpts)
			if err != nil {
				return err
			}
			t, err := buildTable(msg, opts.(*model.TableOptions))
			if err != nil {
				return err
			}
			tables = append(tables, t)
		}
	}

	temp := template.Must(template.New("inline").Parse(tpl))
	temp.Execute(os.Stderr, tables)
	return nil
}

func getRequest() (*plugin.CodeGeneratorRequest, error) {
	data, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		return nil, err
	}

	req := new(plugin.CodeGeneratorRequest)
	if err := proto.Unmarshal(data, req); err != nil {
		return nil, err
	}

	return req, nil
}

func buildTable(msg *descriptor.DescriptorProto, opts *model.TableOptions) (table, error) {
	t := table{
		msg: msg,
	}

	log.Fatal(*msg)

	if opts.Name != "" {
		t.Name = opts.Name
	} else {
		t.Name = *msg.Name
	}
	if strings.ContainsAny(t.Name, "\"\\") {
		return t, errors.New("Invalid characters in table name " + t.Name)
	}
	fieldNames := make(map[string]bool, len(msg.Field))
	for _, f := range msg.Field {
		var typ columnType
		switch *f.Type {
		case descriptor.FieldDescriptorProto_TYPE_FIXED32:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_SFIXED32:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_SINT32:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_ENUM:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_INT32:
			typ = int32Column
		case descriptor.FieldDescriptorProto_TYPE_FIXED64:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_SFIXED64:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_SINT64:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_INT64:
			typ = int64Column
		case descriptor.FieldDescriptorProto_TYPE_BOOL:
			typ = int16Column
		case descriptor.FieldDescriptorProto_TYPE_STRING:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_MESSAGE:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_BYTES:
			typ = bytesColumn
		default:
			return t, errors.New("No type for  " + f.Type.String())
		}
		fieldNames[*f.Name] = true
		t.Columns = append(t.Columns, column{
			Name:       *f.Name,
			ColumnType: typ,
			field:      f,
		})
	}
	if !fieldNames["data"] {
		return t, errors.New("Missing data col on table " + t.Name)
	}
	for _, k := range opts.Key {
		if len(k.Col) == 0 {
			return t, errors.New("No cols in key on table " + t.Name)
		}
		for _, c := range k.Col {
			if !fieldNames[c] {
				return t, errors.New("Unknown col on table " + t.Name)
			}
		}
		name := k.Name
		var typ keyType
		switch k.KeyType {
		case model.KeyType_PRIMARY:
			typ = primaryKey
			if k.Name == "" {
				name = "Primary"
			}
		case model.KeyType_UNIQUE:
			typ = uniqueKey
			if k.Name == "" {
				return t, errors.New("Missing name for key on table " + t.Name)
			}
		case model.KeyType_INDEX:
			typ = indexKey
			if k.Name == "" {
				return t, errors.New("Missing name for key on table " + t.Name)
			}
		default:
			return t, errors.New("Unknown key type on table " + t.Name)
		}
		t.Indexes = append(t.Indexes, index{
			Name:    name,
			KeyType: typ,
			Columns: k.Col,
		})
	}

	return t, nil
}

func main() {
	if err := run(); err != nil {
		log.Fatalln(err)
	}
}
<commit_msg>Output Find and Scan String functions<commit_after>package main

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"text/template"

	"github.com/golang/protobuf/proto"
	descriptor "google/protobuf"
	plugin "google/protobuf/compiler"

	"pixur.org/pixur/schema/db/model"
)

const (
	tpl = `
package tables

import (
  "database/sql"

  "github.com/golang/protobuf/proto"

  "pixur.org/pixur/schema/db"
  "pixur.org/pixur/schema"
)

type Job struct {
  Tx *sql.Tx
}

var SqlTables = []string{
{{range .}}
  {{- .SqlString -}}
{{end -}}
}

{{range .}}
  {{.ScanString}}
  {{.FindString}}
{{end}}
`
)

type columnType string

var (
	int16Column columnType = "smallint"
	int32Column columnType = "integer"
	int64Column columnType = "bigint"
	bytesColumn columnType = "bytea"
)

type keyType string

var (
	primaryKey keyType = "PRIMARY KEY"
	uniqueKey  keyType = "UNIQUE"
	indexKey   keyType = "INDEX"
)

type column struct {
	Name       string
	ColumnType columnType
	field      *descriptor.FieldDescriptorProto
}

type index struct {
	Name    string
	KeyType keyType
	Columns []string
}

type table struct {
	Name    string
	GoType  string
	Columns []column
	Indexes []index
	msg     *descriptor.DescriptorProto
}

func (t table) SqlString() string {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, `  "CREATE TABLE \"%s\" (" +`, t.Name)
	buf.WriteRune('\n')
	for _, col := range t.Columns {
		fmt.Fprintf(&buf, `    "\"%s\" %s NOT NULL, " +`, col.Name, col.ColumnType)
		buf.WriteRune('\n')
	}
	var inlineIndexes []index
	var indexes []index
	for _, idx := range t.Indexes {
		if idx.KeyType == indexKey {
			indexes = append(indexes, idx)
		} else {
			inlineIndexes = append(inlineIndexes, idx)
		}
	}

	for i, idx := range inlineIndexes {
		switch idx.KeyType {
		case indexKey:
			continue
		}
		var cols []string
		for _, col := range idx.Columns {
			cols = append(cols, fmt.Sprintf(`\"%s\"`, col))
		}
		last := ", "
		if i == len(inlineIndexes)-1 {
			last = ""
		}
		fmt.Fprintf(&buf, `    "%s(%s)%s" +`, idx.KeyType, strings.Join(cols, ", "), last)
		buf.WriteRune('\n')
	}
	buf.WriteString("  \");\",\n")
	for _, idx := range indexes {
		var cols []string
		for _, col := range idx.Columns {
			cols = append(cols, fmt.Sprintf(`\"%s\"`, col))
		}
		fmt.Fprintf(&buf, `  "CREATE INDEX \"%s\" ON \"%s\" (%s);",`,
			idx.Name, t.Name, strings.Join(cols, ", "))
		buf.WriteRune('\n')
	}
	return buf.String()
}

func (t table) FindString() string {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, `func (j Job) Find%s(opts Opts) ([]%s, error) {`, t.Name, t.GoType)
	buf.WriteRune('\n')
	fmt.Fprintf(&buf, "\tvar pbs []%s\n", t.GoType)
	fmt.Fprintf(&buf, "\t\terr := j.Scan%s(opts, func(data %s) error {\n", t.Name, t.GoType)
	buf.WriteString("\t\tpbs = append(pbs, data)\n")
	buf.WriteString("\t\treturn nil\n")
	buf.WriteString("\t})\n")
	buf.WriteString("if err != nil {\n")
	buf.WriteString("\treturn nil, err\n")
	buf.WriteString("}\n")
	buf.WriteString("return pbs, nil\n")
	buf.WriteString("}")

	return buf.String()
}

func (t table) ScanString() string {
	var buf bytes.Buffer
	fmt.Fprintf(&buf, `func (j Job) Scan%s(opts Opts, cb func(%s) error) error {`, t.Name, t.GoType)
	buf.WriteRune('\n')
	fmt.Fprintf(&buf, "\t"+`return db.Scan(j.Tx, "%s", opts, func(data []byte) error {`, t.Name)
	buf.WriteRune('\n')
	fmt.Fprintf(&buf, "\t\t"+`var pb %s`, t.GoType)
	buf.WriteRune('\n')
	buf.WriteString("\t\tif err := proto.Unmarshal(data, &pb); err != nil {\n")
	buf.WriteString("\t\t\treturn err\n")
	buf.WriteString("\t\t}\n")
	buf.WriteString("\treturn cb(pb)\n")
	buf.WriteString("})\n")
	buf.WriteString("}\n")
	return buf.String()
}

func run(req *plugin.CodeGeneratorRequest) (*plugin.CodeGeneratorResponse, error) {
	if len(req.FileToGenerate) != 1 {
		return nil, errors.New("Can only generate 1 file.")
	}
	var tables []table
	for _, fd := range req.ProtoFile {
		if fd.GetOptions().GetGoPackage() == "" {
			return nil, errors.New("Must have a Go package")
		}
		for _, msg := range fd.MessageType {
			if msg.Options == nil || !proto.HasExtension(msg.Options, model.E_TabOpts) {
				continue
			}
			opts, err := proto.GetExtension(msg.Options, model.E_TabOpts)
			if err != nil {
				return nil, err
			}
			t, err := buildTable(msg, opts.(*model.TableOptions))
			if err != nil {
				return nil, err
			}
			tables = append(tables, t)
		}
	}

	temp := template.Must(template.New("inline").Parse(tpl))
	var buf bytes.Buffer
	temp.Execute(&buf, tables)
	resp := new(plugin.CodeGeneratorResponse)
	resp.File = []*plugin.CodeGeneratorResponse_File{{
		Name:    proto.String(strings.Replace(req.FileToGenerate[0], ".proto", ".tab.go", -1)),
		Content: proto.String(buf.String()),
	}}

	return resp, nil
}

func getRequest() (*plugin.CodeGeneratorRequest, error) {
	data, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		return nil, err
	}

	req := new(plugin.CodeGeneratorRequest)
	if err := proto.Unmarshal(data, req); err != nil {
		return nil, err
	}

	return req, nil
}

func buildTable(msg *descriptor.DescriptorProto, opts *model.TableOptions) (table, error) {
	t := table{
		msg: msg,
	}

	if opts.Name != "" {
		t.Name = opts.Name
	} else {
		t.Name = *msg.Name
	}
	if strings.ContainsAny(t.Name, "\"\\") {
		return t, errors.New("Invalid characters in table name " + t.Name)
	}
	fieldNames := make(map[string]bool, len(msg.Field))
	for _, f := range msg.Field {
		var typ columnType
		switch *f.Type {
		case descriptor.FieldDescriptorProto_TYPE_FIXED32:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_SFIXED32:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_SINT32:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_ENUM:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_INT32:
			typ = int32Column
		case descriptor.FieldDescriptorProto_TYPE_FIXED64:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_SFIXED64:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_SINT64:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_INT64:
			typ = int64Column
		case descriptor.FieldDescriptorProto_TYPE_BOOL:
			typ = int16Column
		case descriptor.FieldDescriptorProto_TYPE_STRING:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_MESSAGE:
			fallthrough
		case descriptor.FieldDescriptorProto_TYPE_BYTES:
			typ = bytesColumn
		default:
			return t, errors.New("No type for  " + f.Type.String())
		}
		fieldNames[*f.Name] = true
		t.Columns = append(t.Columns, column{
			Name:       *f.Name,
			ColumnType: typ,
			field:      f,
		})

		if *f.Name == "data" {
			parts := strings.Split(*f.TypeName, ".")
			t.GoType = "schema." + parts[len(parts)-1]
		}
	}
	if !fieldNames["data"] {
		return t, errors.New("Missing data col on table " + t.Name)
	}
	for _, k := range opts.Key {
		if len(k.Col) == 0 {
			return t, errors.New("No cols in key on table " + t.Name)
		}
		for _, c := range k.Col {
			if !fieldNames[c] {
				return t, errors.New("Unknown col on table " + t.Name)
			}
		}
		name := k.Name
		var typ keyType
		switch k.KeyType {
		case model.KeyType_PRIMARY:
			typ = primaryKey
			if k.Name == "" {
				name = "Primary"
			}
		case model.KeyType_UNIQUE:
			typ = uniqueKey
			if k.Name == "" {
				return t, errors.New("Missing name for key on table " + t.Name)
			}
		case model.KeyType_INDEX:
			typ = indexKey
			if k.Name == "" {
				return t, errors.New("Missing name for key on table " + t.Name)
			}
		default:
			return t, errors.New("Unknown key type on table " + t.Name)
		}
		t.Indexes = append(t.Indexes, index{
			Name:    name,
			KeyType: typ,
			Columns: k.Col,
		})
	}

	return t, nil
}

func main() {
	req, err := getRequest()
	if err != nil {
		log.Fatalln(err)
	}

	resp, err := run(req)
	if err != nil {
		resp = new(plugin.CodeGeneratorResponse)
		resp.Error = proto.String(err.Error())
	}
	if data, err := proto.Marshal(resp); err != nil {
		log.Fatalln(err)
	} else {
		if _, err := os.Stdout.Write(data); err != nil {
			log.Fatalln(err)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2019 Google LLC
//
// 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
//
//      https://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 e2e

import (
	"bytes"
	"flag"
	"fmt"
	"os"
	"runtime"
	"strings"
	"testing"
	"text/template"
	"time"

	"cloud.google.com/go/profiler/proftest"
	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"

	compute "google.golang.org/api/compute/v1"
)

var (
	gcsLocation       = flag.String("gcs_location", "", "GCS location for the agent")
	runID             = strings.Replace(time.Now().Format("2006-01-02-15-04-05.000000-0700"), ".", "-", -1)
	benchFinishString = "busybench finished profiling"
	errorString       = "failed to set up or run the benchmark"
)

const (
	cloudScope       = "https://www.googleapis.com/auth/cloud-platform"
	storageReadScope = "https://www.googleapis.com/auth/devstorage.read_only"
)

const startupTemplate = `
{{- template "prologue" . }}

# Install dependencies.
{{if .InstallPythonVersion}}
# ppa:deadsnakes/ppa contains desired Python versions.
retry add-apt-repository -y ppa:deadsnakes/ppa >/dev/null
{{end}}
# Force IPv4 to prevent long IPv6 timeouts.
# TODO : Validate this solves the issue. Remove if not.
retry apt-get -o Acquire::ForceIPv4=true update >/dev/null
retry apt-get -o Acquire::ForceIPv4=true install -yq git build-essential python3-distutils {{.PythonDev}} {{if .InstallPythonVersion}}{{.InstallPythonVersion}}{{end}} >/dev/ttyS2

# Install Python dependencies.
retry wget -O /tmp/get-pip.py https://bootstrap.pypa.io/get-pip.py >/dev/null
retry {{.PythonCommand}} /tmp/get-pip.py >/dev/null
retry {{.PythonCommand}} -m pip install --upgrade pyasn1 >/dev/null

# Setup pipenv
retry {{.PythonCommand}} -m pip install pipenv > /dev/null
mkdir bench && cd bench
retry pipenv install > /dev/null

# Fetch agent.
mkdir /tmp/agent
retry gsutil cp gs://{{.GCSLocation}}/* /tmp/agent

# Install agent.
pipenv run {{.PythonCommand}} -m pip install "$(find /tmp/agent -name "google-cloud-profiler*")"

# Run bench app.
export BENCH_DIR="$HOME/bench"

cat << EOF > bench.py
import googlecloudprofiler
import sys
import time
import traceback

def python_bench():
  for counter in range(1, 5000):
    pass

def repeat_bench(dur_sec):
  t_end = time.time() + dur_sec
  while time.time() < t_end or dur_sec == 0:
    python_bench()

if __name__ == '__main__':
  if not {{.VersionCheck}}:
    raise EnvironmentError('Python version %s failed to satisfy "{{.VersionCheck}}".' % str(sys.version_info))

  try:
    googlecloudprofiler.start(
      service='{{.Service}}',
      service_version='1.0.0',
      verbose=3)
  except BaseException:
    sys.exit('Failed to start the profiler: %s' % traceback.format_exc())
  repeat_bench(3 * 60)
EOF

# TODO: Stop ignoring exit code SIGALRM when b/133360821 is fixed.
pipenv run {{.PythonCommand}} bench.py || [ "$?" -eq "142" ]

# Indicate to test that script has finished running.
echo "{{.FinishString}}"

{{ template "epilogue" . -}}
`

type testCase struct {
	proftest.InstanceConfig
	name string
	// Python version to install. Empty string means no installation is needed.
	installPythonVersion string
	// Python command name, e.g "python" or "python3".
	pythonCommand string
	// The python-dev package to install, e.g "python-dev" or "python3.5-dev".
	pythonDev string
	// Used in the bench code to check the Python version, e.g
	// "sys.version_info[:2] == (2.7)".
	versionCheck string
	// Maps profile type to function name wanted for that type. Empty function
	// name means the type should not be profiled.
	wantProfiles map[string]string
}

func (tc *testCase) initializeStartUpScript(template *template.Template) error {
	var buf bytes.Buffer
	err := template.Execute(&buf,
		struct {
			Service              string
			GCSLocation          string
			InstallPythonVersion string
			PythonCommand        string
			PythonDev            string
			VersionCheck         string
			FinishString         string
			ErrorString          string
		}{
			Service:              tc.name,
			GCSLocation:          *gcsLocation,
			InstallPythonVersion: tc.installPythonVersion,
			PythonCommand:        tc.pythonCommand,
			PythonDev:            tc.pythonDev,
			VersionCheck:         tc.versionCheck,
			FinishString:         benchFinishString,
			ErrorString:          errorString,
		})
	if err != nil {
		return fmt.Errorf("failed to render startup script for %s: %v", tc.name, err)
	}
	tc.StartupScript = buf.String()
	return nil
}

func TestAgentIntegration(t *testing.T) {
	projectID := os.Getenv("GCLOUD_TESTS_PYTHON_PROJECT_ID")
	if projectID == "" {
		t.Fatalf("Getenv(GCLOUD_TESTS_PYTHON_PROJECT_ID) got empty string")
	}

	zone := os.Getenv("GCLOUD_TESTS_PYTHON_ZONE")
	if zone == "" {
		t.Fatalf("Getenv(GCLOUD_TESTS_PYTHON_ZONE) got empty string")
	}

	if *gcsLocation == "" {
		t.Fatal("gcsLocation flag is not set")
	}

	ctx := context.Background()

	client, err := google.DefaultClient(ctx, cloudScope)
	if err != nil {
		t.Fatalf("failed to get default client: %v", err)
	}

	computeService, err := compute.New(client)
	if err != nil {
		t.Fatalf("failed to initialize compute Service: %v", err)
	}
	template, err := proftest.BaseStartupTmpl.Parse(startupTemplate)
	if err != nil {
		t.Fatalf("failed to parse startup script template: %v", err)
	}

	gceTr := proftest.GCETestRunner{
		TestRunner: proftest.TestRunner{
			Client: client,
		},
		ComputeService: computeService,
	}

	testcases := []testCase{
		// Test GCE Ubuntu default Python 2, should be Python 2.7.
		{
			InstanceConfig: proftest.InstanceConfig{
				ProjectID:    projectID,
				Zone:         zone,
				Name:         fmt.Sprintf("profiler-test-python2-%s", runID),
				MachineType:  "n1-standard-1",
				ImageProject: "ubuntu-os-cloud",
				ImageFamily:  "ubuntu-1804-lts",
				Scopes:       []string{storageReadScope},
			},
			name: fmt.Sprintf("profiler-test-python2-%s-gce", runID),
			wantProfiles: map[string]string{
				"WALL": "repeat_bench",
				// CPU profiling should be disabled on Python 2.
				"CPU": "",
			},
			pythonCommand: "python2.7",
			pythonDev:     "python-dev",
			versionCheck:  "sys.version_info[:2] == (2, 7)",
		},
		// Test GCE Ubuntu default Python 3, should be Python 3.6 or higher.
		{
			InstanceConfig: proftest.InstanceConfig{
				ProjectID:    projectID,
				Zone:         zone,
				Name:         fmt.Sprintf("profiler-test-python3-%s", runID),
				MachineType:  "n1-standard-1",
				ImageProject: "ubuntu-os-cloud",
				ImageFamily:  "ubuntu-1804-lts",
				Scopes:       []string{storageReadScope},
			},
			name: fmt.Sprintf("profiler-test-python3-%s-gce", runID),
			wantProfiles: map[string]string{
				"WALL": "repeat_bench",
				"CPU":  "repeat_bench",
			},
			pythonCommand: "python3",
			pythonDev:     "python3-dev",
			versionCheck:  "sys.version_info[:2] >= (3, 6)",
		},
		// Test Python 3.5.
		{
			InstanceConfig: proftest.InstanceConfig{
				ProjectID:    projectID,
				Zone:         zone,
				Name:         fmt.Sprintf("profiler-test-python35-%s", runID),
				MachineType:  "n1-standard-1",
				ImageProject: "ubuntu-os-cloud",
				ImageFamily:  "ubuntu-1804-lts",
				Scopes:       []string{storageReadScope},
			},
			name: fmt.Sprintf("profiler-test-python35-%s-gce", runID),
			wantProfiles: map[string]string{
				"CPU": "repeat_bench",
				// Wall profiling should be disabled on Python 3.5.
				"WALL": "",
			},
			installPythonVersion: "python3.5",
			pythonCommand:        "python3.5",
			pythonDev:            "python3.5-dev",
			versionCheck:         "sys.version_info[:2] == (3, 5)",
		},
	}

	// Allow test cases to run in parallel.
	runtime.GOMAXPROCS(len(testcases))

	for _, tc := range testcases {
		tc := tc // capture range variable
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			if err := tc.initializeStartUpScript(template); err != nil {
				t.Fatalf("failed to initialize startup script: %v", err)
			}

			gceTr.StartInstance(ctx, &tc.InstanceConfig)
			defer func() {
				if gceTr.DeleteInstance(ctx, &tc.InstanceConfig); err != nil {
					t.Fatalf("failed to delete instance: %v", err)
				}
			}()

			timeoutCtx, cancel := context.WithTimeout(ctx, time.Minute*20)
			defer cancel()
			if err := gceTr.PollForSerialOutput(timeoutCtx, &tc.InstanceConfig, benchFinishString, errorString); err != nil {
				t.Fatal(err)
			}

			timeNow := time.Now()
			endTime := timeNow.Format(time.RFC3339)
			startTime := timeNow.Add(-1 * time.Hour).Format(time.RFC3339)
			for pType, function := range tc.wantProfiles {
				pr, err := gceTr.TestRunner.QueryProfilesWithZone(tc.ProjectID, tc.name, startTime, endTime, pType, zone)
				if function == "" {
					if err == nil {
						t.Errorf("QueryProfilesWithZone(%s, %s, %s, %s, %s, %s) got profile, want no profile", tc.ProjectID, tc.name, startTime, endTime, pType, zone)
					}
					continue
				}

				if err != nil {
					t.Errorf("QueryProfiles(%s, %s, %s, %s, %s) got error: %v", tc.ProjectID, tc.name, startTime, endTime, pType, err)
					continue
				}

				if err := pr.HasFunction(function); err != nil {
					t.Errorf("Function %s not found in profiles of type %s: %v", function, pType, err)
				}
			}
		})
	}
}
<commit_msg>An internal testing change<commit_after>// Copyright 2019 Google LLC
//
// 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
//
//      https://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 e2e

import (
	"bytes"
	"flag"
	"fmt"
	"os"
	"runtime"
	"strings"
	"testing"
	"text/template"
	"time"

	"cloud.google.com/go/profiler/proftest"
	"golang.org/x/net/context"
	"golang.org/x/oauth2/google"

	compute "google.golang.org/api/compute/v1"
)

var (
	gcsLocation       = flag.String("gcs_location", "", "GCS location for the agent")
	runID             = strings.Replace(time.Now().Format("2006-01-02-15-04-05.000000-0700"), ".", "-", -1)
	benchFinishString = "busybench finished profiling"
	errorString       = "failed to set up or run the benchmark"
)

const (
	cloudScope       = "https://www.googleapis.com/auth/cloud-platform"
	storageReadScope = "https://www.googleapis.com/auth/devstorage.read_only"
)

const startupTemplate = `
{{- template "prologue" . }}

# Install dependencies.
{{if .InstallPythonVersion}}
# ppa:deadsnakes/ppa contains desired Python versions.
retry add-apt-repository -y ppa:deadsnakes/ppa >/dev/null
{{end}}
# Force IPv4 to prevent long IPv6 timeouts.
# TODO : Validate this solves the issue. Remove if not.
retry apt-get -o Acquire::ForceIPv4=true update >/dev/null
retry apt-get -o Acquire::ForceIPv4=true install -yq git build-essential python3-distutils {{.PythonDev}} {{if .InstallPythonVersion}}{{.InstallPythonVersion}}{{end}} >/dev/ttyS2

# Install Python dependencies.
retry wget -O /tmp/get-pip.py https://bootstrap.pypa.io/get-pip.py >/dev/null
retry {{.PythonCommand}} /tmp/get-pip.py >/dev/null
retry {{.PythonCommand}} -m pip install --upgrade pyasn1 >/dev/null

# Setup pipenv
retry {{.PythonCommand}} -m pip install pipenv > /dev/null
mkdir bench && cd bench
retry pipenv install > /dev/null

# Fetch agent.
mkdir /tmp/agent
retry gsutil cp gs://{{.GCSLocation}}/* /tmp/agent

# Install agent.
retry pipenv run {{.PythonCommand}} -m pip install "$(find /tmp/agent -name "google-cloud-profiler*")"

# Run bench app.
export BENCH_DIR="$HOME/bench"

cat << EOF > bench.py
import googlecloudprofiler
import sys
import time
import traceback

def python_bench():
  for counter in range(1, 5000):
    pass

def repeat_bench(dur_sec):
  t_end = time.time() + dur_sec
  while time.time() < t_end or dur_sec == 0:
    python_bench()

if __name__ == '__main__':
  if not {{.VersionCheck}}:
    raise EnvironmentError('Python version %s failed to satisfy "{{.VersionCheck}}".' % str(sys.version_info))

  try:
    googlecloudprofiler.start(
      service='{{.Service}}',
      service_version='1.0.0',
      verbose=3)
  except BaseException:
    sys.exit('Failed to start the profiler: %s' % traceback.format_exc())
  repeat_bench(3 * 60)
EOF

# TODO: Stop ignoring exit code SIGALRM when b/133360821 is fixed.
pipenv run {{.PythonCommand}} bench.py || [ "$?" -eq "142" ]

# Indicate to test that script has finished running.
echo "{{.FinishString}}"

{{ template "epilogue" . -}}
`

type testCase struct {
	proftest.InstanceConfig
	name string
	// Python version to install. Empty string means no installation is needed.
	installPythonVersion string
	// Python command name, e.g "python" or "python3".
	pythonCommand string
	// The python-dev package to install, e.g "python-dev" or "python3.5-dev".
	pythonDev string
	// Used in the bench code to check the Python version, e.g
	// "sys.version_info[:2] == (2.7)".
	versionCheck string
	// Maps profile type to function name wanted for that type. Empty function
	// name means the type should not be profiled.
	wantProfiles map[string]string
}

func (tc *testCase) initializeStartUpScript(template *template.Template) error {
	var buf bytes.Buffer
	err := template.Execute(&buf,
		struct {
			Service              string
			GCSLocation          string
			InstallPythonVersion string
			PythonCommand        string
			PythonDev            string
			VersionCheck         string
			FinishString         string
			ErrorString          string
		}{
			Service:              tc.name,
			GCSLocation:          *gcsLocation,
			InstallPythonVersion: tc.installPythonVersion,
			PythonCommand:        tc.pythonCommand,
			PythonDev:            tc.pythonDev,
			VersionCheck:         tc.versionCheck,
			FinishString:         benchFinishString,
			ErrorString:          errorString,
		})
	if err != nil {
		return fmt.Errorf("failed to render startup script for %s: %v", tc.name, err)
	}
	tc.StartupScript = buf.String()
	return nil
}

func TestAgentIntegration(t *testing.T) {
	projectID := os.Getenv("GCLOUD_TESTS_PYTHON_PROJECT_ID")
	if projectID == "" {
		t.Fatalf("Getenv(GCLOUD_TESTS_PYTHON_PROJECT_ID) got empty string")
	}

	zone := os.Getenv("GCLOUD_TESTS_PYTHON_ZONE")
	if zone == "" {
		t.Fatalf("Getenv(GCLOUD_TESTS_PYTHON_ZONE) got empty string")
	}

	if *gcsLocation == "" {
		t.Fatal("gcsLocation flag is not set")
	}

	ctx := context.Background()

	client, err := google.DefaultClient(ctx, cloudScope)
	if err != nil {
		t.Fatalf("failed to get default client: %v", err)
	}

	computeService, err := compute.New(client)
	if err != nil {
		t.Fatalf("failed to initialize compute Service: %v", err)
	}
	template, err := proftest.BaseStartupTmpl.Parse(startupTemplate)
	if err != nil {
		t.Fatalf("failed to parse startup script template: %v", err)
	}

	gceTr := proftest.GCETestRunner{
		TestRunner: proftest.TestRunner{
			Client: client,
		},
		ComputeService: computeService,
	}

	testcases := []testCase{
		// Test GCE Ubuntu default Python 2, should be Python 2.7.
		{
			InstanceConfig: proftest.InstanceConfig{
				ProjectID:    projectID,
				Zone:         zone,
				Name:         fmt.Sprintf("profiler-test-python2-%s", runID),
				MachineType:  "n1-standard-1",
				ImageProject: "ubuntu-os-cloud",
				ImageFamily:  "ubuntu-1804-lts",
				Scopes:       []string{storageReadScope},
			},
			name: fmt.Sprintf("profiler-test-python2-%s-gce", runID),
			wantProfiles: map[string]string{
				"WALL": "repeat_bench",
				// CPU profiling should be disabled on Python 2.
				"CPU": "",
			},
			pythonCommand: "python2.7",
			pythonDev:     "python-dev",
			versionCheck:  "sys.version_info[:2] == (2, 7)",
		},
		// Test GCE Ubuntu default Python 3, should be Python 3.6 or higher.
		{
			InstanceConfig: proftest.InstanceConfig{
				ProjectID:    projectID,
				Zone:         zone,
				Name:         fmt.Sprintf("profiler-test-python3-%s", runID),
				MachineType:  "n1-standard-1",
				ImageProject: "ubuntu-os-cloud",
				ImageFamily:  "ubuntu-1804-lts",
				Scopes:       []string{storageReadScope},
			},
			name: fmt.Sprintf("profiler-test-python3-%s-gce", runID),
			wantProfiles: map[string]string{
				"WALL": "repeat_bench",
				"CPU":  "repeat_bench",
			},
			pythonCommand: "python3",
			pythonDev:     "python3-dev",
			versionCheck:  "sys.version_info[:2] >= (3, 6)",
		},
		// Test Python 3.5.
		{
			InstanceConfig: proftest.InstanceConfig{
				ProjectID:    projectID,
				Zone:         zone,
				Name:         fmt.Sprintf("profiler-test-python35-%s", runID),
				MachineType:  "n1-standard-1",
				ImageProject: "ubuntu-os-cloud",
				ImageFamily:  "ubuntu-1804-lts",
				Scopes:       []string{storageReadScope},
			},
			name: fmt.Sprintf("profiler-test-python35-%s-gce", runID),
			wantProfiles: map[string]string{
				"CPU": "repeat_bench",
				// Wall profiling should be disabled on Python 3.5.
				"WALL": "",
			},
			installPythonVersion: "python3.5",
			pythonCommand:        "python3.5",
			pythonDev:            "python3.5-dev",
			versionCheck:         "sys.version_info[:2] == (3, 5)",
		},
	}

	// Allow test cases to run in parallel.
	runtime.GOMAXPROCS(len(testcases))

	for _, tc := range testcases {
		tc := tc // capture range variable
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			if err := tc.initializeStartUpScript(template); err != nil {
				t.Fatalf("failed to initialize startup script: %v", err)
			}

			gceTr.StartInstance(ctx, &tc.InstanceConfig)
			defer func() {
				if gceTr.DeleteInstance(ctx, &tc.InstanceConfig); err != nil {
					t.Fatalf("failed to delete instance: %v", err)
				}
			}()

			timeoutCtx, cancel := context.WithTimeout(ctx, time.Minute*20)
			defer cancel()
			if err := gceTr.PollForSerialOutput(timeoutCtx, &tc.InstanceConfig, benchFinishString, errorString); err != nil {
				t.Fatal(err)
			}

			timeNow := time.Now()
			endTime := timeNow.Format(time.RFC3339)
			startTime := timeNow.Add(-1 * time.Hour).Format(time.RFC3339)
			for pType, function := range tc.wantProfiles {
				pr, err := gceTr.TestRunner.QueryProfilesWithZone(tc.ProjectID, tc.name, startTime, endTime, pType, zone)
				if function == "" {
					if err == nil {
						t.Errorf("QueryProfilesWithZone(%s, %s, %s, %s, %s, %s) got profile, want no profile", tc.ProjectID, tc.name, startTime, endTime, pType, zone)
					}
					continue
				}

				if err != nil {
					t.Errorf("QueryProfiles(%s, %s, %s, %s, %s) got error: %v", tc.ProjectID, tc.name, startTime, endTime, pType, err)
					continue
				}

				if err := pr.HasFunction(function); err != nil {
					t.Errorf("Function %s not found in profiles of type %s: %v", function, pType, err)
				}
			}
		})
	}
}
<|endoftext|>