<commit_before>package commands

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"strings"
	"sync"

	"github.com/git-lfs/git-lfs/errors"
	"github.com/git-lfs/git-lfs/filepathfilter"
	"github.com/git-lfs/git-lfs/git"
	"github.com/git-lfs/git-lfs/lfs"
	"github.com/git-lfs/git-lfs/tq"
	"github.com/spf13/cobra"
)

const (
	// cleanFilterBufferCapacity is the desired capacity of the
	// `*git.PacketWriter`'s internal buffer when the filter protocol
	// dictates the "clean" command. 512 bytes is (in most cases) enough to
	// hold an entire LFS pointer in memory.
	cleanFilterBufferCapacity = 512

	// smudgeFilterBufferCapacity is the desired capacity of the
	// `*git.PacketWriter`'s internal buffer when the filter protocol
	// dictates the "smudge" command.
	smudgeFilterBufferCapacity = git.MaxPacketLength
)

// filterSmudgeSkip is a command-line flag owned by the `filter-process` command
// dictating whether or not to skip the smudging process, leaving pointers as-is
// in the working tree.
var filterSmudgeSkip bool

func filterCommand(cmd *cobra.Command, args []string) {
	requireStdin("This command should be run by the Git filter process")
	lfs.InstallHooks(false)

	s := git.NewFilterProcessScanner(os.Stdin, os.Stdout)

	if err := s.Init(); err != nil {
		ExitWithError(err)
	}

	caps, err := s.NegotiateCapabilities()
	if err != nil {
		ExitWithError(err)
	}

	var supportsDelay bool
	for _, cap := range caps {
		if cap == "capability=delay" {
			supportsDelay = true
			break
		}
	}

	skip := filterSmudgeSkip || cfg.Os.Bool("GIT_LFS_SKIP_SMUDGE", false)
	filter := filepathfilter.New(cfg.FetchIncludePaths(), cfg.FetchExcludePaths())

	ptrs := make(map[string]*lfs.Pointer)

	var q *tq.TransferQueue
	closeOnce := new(sync.Once)
	available := make(chan *tq.Transfer)

	if supportsDelay {
		q = tq.NewTransferQueue(tq.Download, getTransferManifest(), cfg.CurrentRemote)
		go infiniteTransferBuffer(q, available)
	}

	var malformed []string
	var malformedOnWindows []string
	for s.Scan() {
		var n int64
		var err error
		var delayed bool
		var w *git.PktlineWriter

		req := s.Request()

		if !(req.Header["command"] == "smudge" && req.Header["can-delay"] == "1") && !(req.Header["command"] == "list_available_blobs") {
			s.WriteStatus(statusFromErr(nil))
		}

		switch req.Header["command"] {
		case "clean":
			w = git.NewPktlineWriter(os.Stdout, cleanFilterBufferCapacity)

			var ptr *lfs.Pointer
			ptr, err = clean(w, req.Payload, req.Header["pathname"], -1)

			if ptr != nil {
				n = ptr.Size
			}
		case "smudge":
			w = git.NewPktlineWriter(os.Stdout, smudgeFilterBufferCapacity)
			if req.Header["can-delay"] == "1" {
				var ptr *lfs.Pointer

				n, delayed, ptr, err = delayedSmudge(s, w, req.Payload, q, req.Header["pathname"], skip, filter)

				if delayed {
					ptrs[req.Header["pathname"]] = ptr
				}
			} else {
				from, ferr := incomingOrCached(req.Payload, ptrs[req.Header["pathname"]])
				if ferr != nil {
					break
				}

				n, err = smudge(w, from, req.Header["pathname"], skip, filter)
				if err == nil {
					delete(ptrs, req.Header["pathname"])
				}
			}
		case "list_available_blobs":
			closeOnce.Do(func() {
				// The first time that Git sends us the
				// 'list_available_blobs' command, it is given
				// that no more smudge commands will be issued
				// with _new_ checkout entries.
				//
				// This means that, by the time that we're here,
				// we have seen all entries in the checkout, and
				// should therefore instruct the transfer queue
				// to make a batch out of whatever remaining
				// items it has, and then close itself.
				//
				// This function call is wrapped in a
				// `sync.(*Once).Do()` call so we only call
				// `q.Wait()` once, and is called via a
				// goroutine since `q.Wait()` is blocking.
				go q.Wait()
			})

			// The first, and all subsequent calls to
			// list_available_blobs, we read items from `tq.Watch()`
			// until a read from that channel becomes blocking (in
			// other words, we read until there are no more items
			// immediately ready to be sent back to Git).
			paths := pathnames(readAvailable(available))
			if len(paths) == 0 {
				// If `len(paths) == 0`, `tq.Watch()` has
				// closed, indicating that all items have been
				// completely processed, and therefore, sent
				// back to Git for checkout.
				for path, _ := range ptrs {
					// If we sent a path to Git but it
					// didn't ask for the smudge contents,
					// that path is available and Git should
					// accept it later.
					paths = append(paths, fmt.Sprintf("pathname=%s", path))
				}
			}
			err = s.WriteList(paths)
		default:
			ExitWithError(fmt.Errorf("Unknown command %q", req.Header["command"]))
		}

		if errors.IsNotAPointerError(err) {
			malformed = append(malformed, req.Header["pathname"])
			err = nil
		} else if possiblyMalformedObjectSize(n) {
			malformedOnWindows = append(malformedOnWindows, req.Header["pathname"])
		}

		var status string
		if delayed {
			// If delayed, there is no need to call w.Flush() since
			// no data was written. Calculate the status from the
			// given error using 'delayedStatusFromErr'.
			status = delayedStatusFromErr(err)
		} else if ferr := w.Flush(); ferr != nil {
			// Otherwise, we do need to call w.Flush(), since we
			// have to assume that data was written. If the flush
			// operation was unsuccessful, calculate the status
			// using 'statusFromErr'.
			status = statusFromErr(ferr)
		} else {
			// If the above flush was successful, we calculate the
			// status from the above clean, smudge, or
			// list_available_blobs command using statusFromErr,
			// since we did not delay.
			status = statusFromErr(err)
		}

		s.WriteStatus(status)
	}

	if len(malformed) > 0 {
		fmt.Fprintf(os.Stderr, "Encountered %d file(s) that should have been pointers, but weren't:\n", len(malformed))
		for _, m := range malformed {
			fmt.Fprintf(os.Stderr, "\t%s\n", m)
		}
	}

	if len(malformedOnWindows) > 0 {
		fmt.Fprintf(os.Stderr, "Encountered %d file(s) that may not have been copied correctly on Windows:\n")

		for _, m := range malformedOnWindows {
			fmt.Fprintf(os.Stderr, "\t%s\n", m)
		}

		fmt.Fprintf(os.Stderr, "\nSee: `git lfs help smudge` for more details.\n")
	}

	if err := s.Err(); err != nil && err != io.EOF {
		ExitWithError(err)
	}
}

// infiniteTransferBuffer streams the results of q.Watch() into "available" as
// if available had an infinite channel buffer.
func infiniteTransferBuffer(q *tq.TransferQueue, available chan<- *tq.Transfer) {
	// Stream results from q.Watch() into chan "available" via an infinite
	// buffer.

	watch := q.Watch()

	// pending is used to keep track of an ordered list of available
	// `*tq.Transfer`'s that cannot be written to "available" without
	// blocking.
	var pending []*tq.Transfer

	for {
		if len(pending) > 0 {
			select {
			case t, ok := <-watch:
				if !ok {
					// If the list of pending elements is
					// non-empty, stream them out (even if
					// they block), and then close().
					for _, t = range pending {
						available <- t
					}
					close(available)
					return
				}
				pending = append(pending, t)
			case available <- pending[0]:
				// Otherwise, dequeue and shift the first
				// element from pending onto available.
				pending = pending[1:]
			}
		} else {
			t, ok := <-watch
			if !ok {
				// If watch is closed, the "tq" is done, and
				// there are no items on the buffer.  Return
				// immediately.
				close(available)
				return
			}

			select {
			case available <- t:
			// Copy an item directly from <-watch onto available<-.
			default:
				// Otherwise, if that would have blocked, make
				// the new read pending.
				pending = append(pending, t)
			}
		}
	}
}

// incomingOrCached returns an io.Reader that is either the contents of the
// given io.Reader "r", or the encoded contents of "ptr". It returns an error if
// there was an error reading from "r".
//
// This is done because when a `command=smudge` with `can-delay=0` is issued,
// the entry's contents are not sent, and must be re-encoded from the stored
// pointer corresponding to the request's filepath.
func incomingOrCached(r io.Reader, ptr *lfs.Pointer) (io.Reader, error) {
	buf := make([]byte, 1024)
	n, err := r.Read(buf)
	buf = buf[:n]

	if n == 0 {
		if ptr == nil {
			// If we read no data from the given io.Reader "r" _and_
			// there was no data to fall back on, return an empty
			// io.Reader yielding no data.
			return bytes.NewReader(buf), nil
		}
		// If we read no data from the given io.Reader "r", _and_ there
		// is a pointer that we can fall back on, return an io.Reader
		// that yields the encoded version of the given pointer.
		return strings.NewReader(ptr.Encoded()), nil
	}

	if err == io.EOF {
		return bytes.NewReader(buf), nil
	}
	return io.MultiReader(bytes.NewReader(buf), r), err
}

// readAvailable satisfies the accumulation semantics for the
// 'list_available_blobs' command. It accumulates items until:
//
// 1. Reading from the channel of available items blocks, or ...
// 2. There is one item available, or ...
// 3. The 'tq.TransferQueue' is completed.
func readAvailable(ch <-chan *tq.Transfer) []*tq.Transfer {
	ts := make([]*tq.Transfer, 0, 100)

	for {
		select {
		case t, ok := <-ch:
			if !ok {
				return ts
			}
			ts = append(ts, t)
		default:
			if len(ts) > 0 {
				return ts
			}

			t, ok := <-ch
			if !ok {
				return ts
			}
			return append(ts, t)
		}
	}

	return ts
}

// pathnames formats a list of *tq.Transfers as a valid response to the
// 'list_available_blobs' command.
func pathnames(ts []*tq.Transfer) []string {
	pathnames := make([]string, 0, len(ts))
	for _, t := range ts {
		pathnames = append(pathnames, fmt.Sprintf("pathname=%s", t.Name))
	}

	return pathnames
}

// statusFromErr returns the status code that should be sent over the filter
// protocol based on a given error, "err".
func statusFromErr(err error) string {
	if err != nil && err != io.EOF {
		return "error"
	}
	return "success"
}

// delayedStatusFromErr returns the status code that should be sent over the
// filter protocol based on a given error, "err" when the blob smudge operation
// was delayed.
func delayedStatusFromErr(err error) string {
	if err != nil && err != io.EOF {
		return "error"
	}
	return "delayed"
}

func init() {
	RegisterCommand("filter-process", filterCommand, func(cmd *cobra.Command) {
		cmd.Flags().BoolVarP(&filterSmudgeSkip, "skip", "s", false, "")
	})
}
<commit_msg>commands: call 's.WriteStatus' explicitly<commit_after>package commands

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"strings"
	"sync"

	"github.com/git-lfs/git-lfs/errors"
	"github.com/git-lfs/git-lfs/filepathfilter"
	"github.com/git-lfs/git-lfs/git"
	"github.com/git-lfs/git-lfs/lfs"
	"github.com/git-lfs/git-lfs/tq"
	"github.com/spf13/cobra"
)

const (
	// cleanFilterBufferCapacity is the desired capacity of the
	// `*git.PacketWriter`'s internal buffer when the filter protocol
	// dictates the "clean" command. 512 bytes is (in most cases) enough to
	// hold an entire LFS pointer in memory.
	cleanFilterBufferCapacity = 512

	// smudgeFilterBufferCapacity is the desired capacity of the
	// `*git.PacketWriter`'s internal buffer when the filter protocol
	// dictates the "smudge" command.
	smudgeFilterBufferCapacity = git.MaxPacketLength
)

// filterSmudgeSkip is a command-line flag owned by the `filter-process` command
// dictating whether or not to skip the smudging process, leaving pointers as-is
// in the working tree.
var filterSmudgeSkip bool

func filterCommand(cmd *cobra.Command, args []string) {
	requireStdin("This command should be run by the Git filter process")
	lfs.InstallHooks(false)

	s := git.NewFilterProcessScanner(os.Stdin, os.Stdout)

	if err := s.Init(); err != nil {
		ExitWithError(err)
	}

	caps, err := s.NegotiateCapabilities()
	if err != nil {
		ExitWithError(err)
	}

	var supportsDelay bool
	for _, cap := range caps {
		if cap == "capability=delay" {
			supportsDelay = true
			break
		}
	}

	skip := filterSmudgeSkip || cfg.Os.Bool("GIT_LFS_SKIP_SMUDGE", false)
	filter := filepathfilter.New(cfg.FetchIncludePaths(), cfg.FetchExcludePaths())

	ptrs := make(map[string]*lfs.Pointer)

	var q *tq.TransferQueue
	closeOnce := new(sync.Once)
	available := make(chan *tq.Transfer)

	if supportsDelay {
		q = tq.NewTransferQueue(tq.Download, getTransferManifest(), cfg.CurrentRemote)
		go infiniteTransferBuffer(q, available)
	}

	var malformed []string
	var malformedOnWindows []string
	for s.Scan() {
		var n int64
		var err error
		var delayed bool
		var w *git.PktlineWriter

		req := s.Request()

		switch req.Header["command"] {
		case "clean":
			s.WriteStatus(statusFromErr(nil))
			w = git.NewPktlineWriter(os.Stdout, cleanFilterBufferCapacity)

			var ptr *lfs.Pointer
			ptr, err = clean(w, req.Payload, req.Header["pathname"], -1)

			if ptr != nil {
				n = ptr.Size
			}
		case "smudge":
			w = git.NewPktlineWriter(os.Stdout, smudgeFilterBufferCapacity)
			if req.Header["can-delay"] == "1" {
				var ptr *lfs.Pointer

				n, delayed, ptr, err = delayedSmudge(s, w, req.Payload, q, req.Header["pathname"], skip, filter)

				if delayed {
					ptrs[req.Header["pathname"]] = ptr
				}
			} else {
				s.WriteStatus(statusFromErr(nil))
				from, ferr := incomingOrCached(req.Payload, ptrs[req.Header["pathname"]])
				if ferr != nil {
					break
				}

				n, err = smudge(w, from, req.Header["pathname"], skip, filter)
				if err == nil {
					delete(ptrs, req.Header["pathname"])
				}
			}
		case "list_available_blobs":
			closeOnce.Do(func() {
				// The first time that Git sends us the
				// 'list_available_blobs' command, it is given
				// that no more smudge commands will be issued
				// with _new_ checkout entries.
				//
				// This means that, by the time that we're here,
				// we have seen all entries in the checkout, and
				// should therefore instruct the transfer queue
				// to make a batch out of whatever remaining
				// items it has, and then close itself.
				//
				// This function call is wrapped in a
				// `sync.(*Once).Do()` call so we only call
				// `q.Wait()` once, and is called via a
				// goroutine since `q.Wait()` is blocking.
				go q.Wait()
			})

			// The first, and all subsequent calls to
			// list_available_blobs, we read items from `tq.Watch()`
			// until a read from that channel becomes blocking (in
			// other words, we read until there are no more items
			// immediately ready to be sent back to Git).
			paths := pathnames(readAvailable(available))
			if len(paths) == 0 {
				// If `len(paths) == 0`, `tq.Watch()` has
				// closed, indicating that all items have been
				// completely processed, and therefore, sent
				// back to Git for checkout.
				for path, _ := range ptrs {
					// If we sent a path to Git but it
					// didn't ask for the smudge contents,
					// that path is available and Git should
					// accept it later.
					paths = append(paths, fmt.Sprintf("pathname=%s", path))
				}
			}
			err = s.WriteList(paths)
		default:
			ExitWithError(fmt.Errorf("Unknown command %q", req.Header["command"]))
		}

		if errors.IsNotAPointerError(err) {
			malformed = append(malformed, req.Header["pathname"])
			err = nil
		} else if possiblyMalformedObjectSize(n) {
			malformedOnWindows = append(malformedOnWindows, req.Header["pathname"])
		}

		var status string
		if delayed {
			// If delayed, there is no need to call w.Flush() since
			// no data was written. Calculate the status from the
			// given error using 'delayedStatusFromErr'.
			status = delayedStatusFromErr(err)
		} else if ferr := w.Flush(); ferr != nil {
			// Otherwise, we do need to call w.Flush(), since we
			// have to assume that data was written. If the flush
			// operation was unsuccessful, calculate the status
			// using 'statusFromErr'.
			status = statusFromErr(ferr)
		} else {
			// If the above flush was successful, we calculate the
			// status from the above clean, smudge, or
			// list_available_blobs command using statusFromErr,
			// since we did not delay.
			status = statusFromErr(err)
		}

		s.WriteStatus(status)
	}

	if len(malformed) > 0 {
		fmt.Fprintf(os.Stderr, "Encountered %d file(s) that should have been pointers, but weren't:\n", len(malformed))
		for _, m := range malformed {
			fmt.Fprintf(os.Stderr, "\t%s\n", m)
		}
	}

	if len(malformedOnWindows) > 0 {
		fmt.Fprintf(os.Stderr, "Encountered %d file(s) that may not have been copied correctly on Windows:\n")

		for _, m := range malformedOnWindows {
			fmt.Fprintf(os.Stderr, "\t%s\n", m)
		}

		fmt.Fprintf(os.Stderr, "\nSee: `git lfs help smudge` for more details.\n")
	}

	if err := s.Err(); err != nil && err != io.EOF {
		ExitWithError(err)
	}
}

// infiniteTransferBuffer streams the results of q.Watch() into "available" as
// if available had an infinite channel buffer.
func infiniteTransferBuffer(q *tq.TransferQueue, available chan<- *tq.Transfer) {
	// Stream results from q.Watch() into chan "available" via an infinite
	// buffer.

	watch := q.Watch()

	// pending is used to keep track of an ordered list of available
	// `*tq.Transfer`'s that cannot be written to "available" without
	// blocking.
	var pending []*tq.Transfer

	for {
		if len(pending) > 0 {
			select {
			case t, ok := <-watch:
				if !ok {
					// If the list of pending elements is
					// non-empty, stream them out (even if
					// they block), and then close().
					for _, t = range pending {
						available <- t
					}
					close(available)
					return
				}
				pending = append(pending, t)
			case available <- pending[0]:
				// Otherwise, dequeue and shift the first
				// element from pending onto available.
				pending = pending[1:]
			}
		} else {
			t, ok := <-watch
			if !ok {
				// If watch is closed, the "tq" is done, and
				// there are no items on the buffer.  Return
				// immediately.
				close(available)
				return
			}

			select {
			case available <- t:
			// Copy an item directly from <-watch onto available<-.
			default:
				// Otherwise, if that would have blocked, make
				// the new read pending.
				pending = append(pending, t)
			}
		}
	}
}

// incomingOrCached returns an io.Reader that is either the contents of the
// given io.Reader "r", or the encoded contents of "ptr". It returns an error if
// there was an error reading from "r".
//
// This is done because when a `command=smudge` with `can-delay=0` is issued,
// the entry's contents are not sent, and must be re-encoded from the stored
// pointer corresponding to the request's filepath.
func incomingOrCached(r io.Reader, ptr *lfs.Pointer) (io.Reader, error) {
	buf := make([]byte, 1024)
	n, err := r.Read(buf)
	buf = buf[:n]

	if n == 0 {
		if ptr == nil {
			// If we read no data from the given io.Reader "r" _and_
			// there was no data to fall back on, return an empty
			// io.Reader yielding no data.
			return bytes.NewReader(buf), nil
		}
		// If we read no data from the given io.Reader "r", _and_ there
		// is a pointer that we can fall back on, return an io.Reader
		// that yields the encoded version of the given pointer.
		return strings.NewReader(ptr.Encoded()), nil
	}

	if err == io.EOF {
		return bytes.NewReader(buf), nil
	}
	return io.MultiReader(bytes.NewReader(buf), r), err
}

// readAvailable satisfies the accumulation semantics for the
// 'list_available_blobs' command. It accumulates items until:
//
// 1. Reading from the channel of available items blocks, or ...
// 2. There is one item available, or ...
// 3. The 'tq.TransferQueue' is completed.
func readAvailable(ch <-chan *tq.Transfer) []*tq.Transfer {
	ts := make([]*tq.Transfer, 0, 100)

	for {
		select {
		case t, ok := <-ch:
			if !ok {
				return ts
			}
			ts = append(ts, t)
		default:
			if len(ts) > 0 {
				return ts
			}

			t, ok := <-ch
			if !ok {
				return ts
			}
			return append(ts, t)
		}
	}

	return ts
}

// pathnames formats a list of *tq.Transfers as a valid response to the
// 'list_available_blobs' command.
func pathnames(ts []*tq.Transfer) []string {
	pathnames := make([]string, 0, len(ts))
	for _, t := range ts {
		pathnames = append(pathnames, fmt.Sprintf("pathname=%s", t.Name))
	}

	return pathnames
}

// statusFromErr returns the status code that should be sent over the filter
// protocol based on a given error, "err".
func statusFromErr(err error) string {
	if err != nil && err != io.EOF {
		return "error"
	}
	return "success"
}

// delayedStatusFromErr returns the status code that should be sent over the
// filter protocol based on a given error, "err" when the blob smudge operation
// was delayed.
func delayedStatusFromErr(err error) string {
	if err != nil && err != io.EOF {
		return "error"
	}
	return "delayed"
}

func init() {
	RegisterCommand("filter-process", filterCommand, func(cmd *cobra.Command) {
		cmd.Flags().BoolVarP(&filterSmudgeSkip, "skip", "s", false, "")
	})
}
<|endoftext|>


<commit_before>package container551

import (
	"github.com/go51/auth551"
	"github.com/go51/cookie551"
	"github.com/go51/log551"
	"github.com/go51/memcache551"
	"github.com/go51/model551"
	"github.com/go51/mysql551"
	"github.com/go51/repository551"
	"github.com/go51/secure551"
	"github.com/go51/string551"
	xoauth2 "golang.org/x/oauth2"
	"net/http"
	"strconv"
)

type Container struct {
	sid     string
	ssid    string
	w       http.ResponseWriter
	r       *http.Request
	logger  *log551.Log551
	cookie  *cookie551.Cookie
	db      *mysql551.Mysql
	session *memcache551.Memcache
	model   *model551.Model
	auth    *auth551.Auth
	user    *auth551.UserModel
	options map[string]string
}

func New() *Container {
	return &Container{}
}

func (c *Container) SetSID(sid string) {
	c.sid = sid
	c.ssid = sid[:10]
}

func (c *Container) SID() string {
	return c.sid
}
func (c *Container) SSID() string {
	return c.ssid
}

func (c *Container) SetResponseWriter(w http.ResponseWriter) {
	c.w = w
}

func (c *Container) ResponseWriter() http.ResponseWriter {
	return c.w
}

func (c *Container) SetRequest(r *http.Request) {
	c.r = r
}

func (c *Container) Request() *http.Request {
	return c.r
}

func (c *Container) SetLogger(logger *log551.Log551) {
	c.logger = logger
}

func (c *Container) Logger() *log551.Log551 {
	return c.logger
}

func (c *Container) SetCookie(cookie *cookie551.Cookie) {
	c.cookie = cookie
}

func (c *Container) Cookie() *cookie551.Cookie {
	return c.cookie
}

func (c *Container) SetDb(db *mysql551.Mysql) {
	c.db = db
}

func (c *Container) Db() *mysql551.Mysql {
	return c.db
}

func (c *Container) SetSession(session *memcache551.Memcache) {
	c.session = session
}

func (c *Container) Session() *memcache551.Memcache {
	return c.session
}

func (c *Container) SetModel(modelManager *model551.Model) {
	c.model = modelManager
}

func (c *Container) ModelManager() *model551.Model {
	return c.model
}

func (c *Container) SetAuth(auth *auth551.Auth) {
	c.auth = auth

	if c.user != nil {
		return
	}

	// Load user from session
	c.session.GetModel("reminder_user", &c.user)

	// Get user id from cookie
	id, err := c.getRemindId()
	if err != nil {
		return
	}

	// Get user model from database
	c.user = c.getUser(id)

	// Set user model to session
	c.session.Set("reminder_user", c.user)

	return

}

func (c *Container) getRemindId() (int64, error) {
	cookieId, err := c.cookie.Get(c.auth.CookieKeyName())
	if err != nil {
		return 0, err
	}

	sid := secure551.Decrypted(cookieId, c.auth.MasterKey())
	id, err := strconv.ParseInt(sid, 10, 64)
	if err != nil {
		return 0, err
	}

	return id, nil

}

func (c *Container) getUser(id int64) *auth551.UserModel {
	repo := repository551.Load()
	miUser := c.ModelManager().Get("UserModel")
	mUser := repo.Find(c.db, miUser, id)
	user, ok := mUser.(*auth551.UserModel)
	if !ok {
		return nil
	}

	return user
}

func (c *Container) Auth() *auth551.Auth {
	return c.auth
}

func (c *Container) SignIn(user *auth551.UserModel) {
	// Set remind id to cookie
	id := string551.Right("0000000000000000"+strconv.FormatInt(user.Id, 10), 16)
	secureId := secure551.Encrypted(id, c.auth.MasterKey())
	c.cookie.Set(c.auth.CookieKeyName(), secureId, 60*60*24*365)

	// Set user model to session
	c.session.Set("reminder_user", user)

}

func (c *Container) Logout() {
	c.cookie.Delete(c.auth.CookieKeyName())
	c.session.Delete("user")
}

func (c *Container) IsLogin() bool {
	return c.user != nil
}

func (c *Container) User() *auth551.UserModel {
	return c.user
}

func (c *Container) ApiClient(vendor auth551.AuthVendor) *http.Client {
	if c.user == nil {
		return nil
	}
	repo := repository551.Load()
	miUserToken := c.model.Get("UserTokenModel")
	param := map[string]interface{}{
		"user_id": c.user.Id,
	}
	mUserToken := repo.FindOneBy(c.db, miUserToken, param)
	userToken, ok := mUserToken.(*auth551.UserTokenModel)
	if !ok {
		return nil
	}

	token := &xoauth2.Token{
		AccessToken:  userToken.AccessToken,
		TokenType:    userToken.TokenType,
		RefreshToken: userToken.RefreshToken,
		Expiry:       userToken.Expiry,
	}

	return c.auth.Client(vendor, token)

}

func (c *Container) SetCommandOptions(options map[string]string) {
	c.options = options
}

func (c *Container) CommandOption(name string) string {
	return c.options[name]
}
<commit_msg>refs #27 ログアウト時の削除するセッションキーを修正&メソッド名を変更<commit_after>package container551

import (
	"github.com/go51/auth551"
	"github.com/go51/cookie551"
	"github.com/go51/log551"
	"github.com/go51/memcache551"
	"github.com/go51/model551"
	"github.com/go51/mysql551"
	"github.com/go51/repository551"
	"github.com/go51/secure551"
	"github.com/go51/string551"
	xoauth2 "golang.org/x/oauth2"
	"net/http"
	"strconv"
)

type Container struct {
	sid     string
	ssid    string
	w       http.ResponseWriter
	r       *http.Request
	logger  *log551.Log551
	cookie  *cookie551.Cookie
	db      *mysql551.Mysql
	session *memcache551.Memcache
	model   *model551.Model
	auth    *auth551.Auth
	user    *auth551.UserModel
	options map[string]string
}

func New() *Container {
	return &Container{}
}

func (c *Container) SetSID(sid string) {
	c.sid = sid
	c.ssid = sid[:10]
}

func (c *Container) SID() string {
	return c.sid
}
func (c *Container) SSID() string {
	return c.ssid
}

func (c *Container) SetResponseWriter(w http.ResponseWriter) {
	c.w = w
}

func (c *Container) ResponseWriter() http.ResponseWriter {
	return c.w
}

func (c *Container) SetRequest(r *http.Request) {
	c.r = r
}

func (c *Container) Request() *http.Request {
	return c.r
}

func (c *Container) SetLogger(logger *log551.Log551) {
	c.logger = logger
}

func (c *Container) Logger() *log551.Log551 {
	return c.logger
}

func (c *Container) SetCookie(cookie *cookie551.Cookie) {
	c.cookie = cookie
}

func (c *Container) Cookie() *cookie551.Cookie {
	return c.cookie
}

func (c *Container) SetDb(db *mysql551.Mysql) {
	c.db = db
}

func (c *Container) Db() *mysql551.Mysql {
	return c.db
}

func (c *Container) SetSession(session *memcache551.Memcache) {
	c.session = session
}

func (c *Container) Session() *memcache551.Memcache {
	return c.session
}

func (c *Container) SetModel(modelManager *model551.Model) {
	c.model = modelManager
}

func (c *Container) ModelManager() *model551.Model {
	return c.model
}

func (c *Container) SetAuth(auth *auth551.Auth) {
	c.auth = auth

	if c.user != nil {
		return
	}

	// Load user from session
	c.session.GetModel("reminder_user", &c.user)

	// Get user id from cookie
	id, err := c.getRemindId()
	if err != nil {
		return
	}

	// Get user model from database
	c.user = c.getUser(id)

	// Set user model to session
	c.session.Set("reminder_user", c.user)

	return

}

func (c *Container) getRemindId() (int64, error) {
	cookieId, err := c.cookie.Get(c.auth.CookieKeyName())
	if err != nil {
		return 0, err
	}

	sid := secure551.Decrypted(cookieId, c.auth.MasterKey())
	id, err := strconv.ParseInt(sid, 10, 64)
	if err != nil {
		return 0, err
	}

	return id, nil

}

func (c *Container) getUser(id int64) *auth551.UserModel {
	repo := repository551.Load()
	miUser := c.ModelManager().Get("UserModel")
	mUser := repo.Find(c.db, miUser, id)
	user, ok := mUser.(*auth551.UserModel)
	if !ok {
		return nil
	}

	return user
}

func (c *Container) Auth() *auth551.Auth {
	return c.auth
}

func (c *Container) SignIn(user *auth551.UserModel) {
	// Set remind id to cookie
	id := string551.Right("0000000000000000"+strconv.FormatInt(user.Id, 10), 16)
	secureId := secure551.Encrypted(id, c.auth.MasterKey())
	c.cookie.Set(c.auth.CookieKeyName(), secureId, 60*60*24*365)

	// Set user model to session
	c.session.Set("reminder_user", user)

}

func (c *Container) SignOut() {
	c.cookie.Delete(c.auth.CookieKeyName())
	c.session.Delete("reminder_user")
}

func (c *Container) IsSignIn() bool {
	return c.user != nil
}

func (c *Container) User() *auth551.UserModel {
	return c.user
}

func (c *Container) ApiClient(vendor auth551.AuthVendor) *http.Client {
	if c.user == nil {
		return nil
	}
	repo := repository551.Load()
	miUserToken := c.model.Get("UserTokenModel")
	param := map[string]interface{}{
		"user_id": c.user.Id,
	}
	mUserToken := repo.FindOneBy(c.db, miUserToken, param)
	userToken, ok := mUserToken.(*auth551.UserTokenModel)
	if !ok {
		return nil
	}

	token := &xoauth2.Token{
		AccessToken:  userToken.AccessToken,
		TokenType:    userToken.TokenType,
		RefreshToken: userToken.RefreshToken,
		Expiry:       userToken.Expiry,
	}

	return c.auth.Client(vendor, token)

}

func (c *Container) SetCommandOptions(options map[string]string) {
	c.options = options
}

func (c *Container) CommandOption(name string) string {
	return c.options[name]
}
<|endoftext|>


<commit_before>// +build linux

package tcp_test

import (
	"context"
	"testing"

	"v2ray.com/core/testing/assert"
	"v2ray.com/core/testing/servers/tcp"
)

func TestGetOriginalDestination(t *testing.T) {
	assert := assert.On(t)

	tcpServer := tcp.Server{
		MsgProcessor: xor,
	}
	dest, err := tcpServer.Start()
	assert.Error(err).IsNil()
	defer tcpServer.Close()

	conn, err := Dial(context.Background(), dest)
	assert.Error(err).IsNil()

	_, err := GetOriginalDestination(conn)
	assert.String(err.Error()).Contains("failed to call getsockopt")
}
<commit_msg>fix test break<commit_after>// +build linux

package tcp_test

import (
	"context"
	"testing"

	"v2ray.com/core/testing/assert"
	"v2ray.com/core/testing/servers/tcp"
	. "v2ray.com/core/transport/internet/tcp"
)

func TestGetOriginalDestination(t *testing.T) {
	assert := assert.On(t)

	tcpServer := tcp.Server{}
	dest, err := tcpServer.Start()
	assert.Error(err).IsNil()
	defer tcpServer.Close()

	conn, err := Dial(context.Background(), dest)
	assert.Error(err).IsNil()
	defer conn.Close()

	_, err := GetOriginalDestination(conn)
	assert.String(err.Error()).Contains("failed to call getsockopt")
}
<|endoftext|>


<commit_before>package components

import (
	"fmt"
	"os"
	"testing"

	"log"

	"github.com/scipipe/scipipe"
)

func TestFileCombinator(t *testing.T) {
	letters := []string{"a", "b"}
	numbers := []string{"1", "2", "3"}

	// Create letter files
	for _, s := range letters {
		fName := "/tmp/letterfile_" + s + ".txt"
		f, err := os.Create(fName)
		if err != nil {
			log.Fatalf("File could not be created: %s\n", fName)
		}
		f.WriteString(s)
	}

	// Create number files
	for _, s := range numbers {
		fName := "/tmp/numberfile_" + s + ".txt"
		f, err := os.Create(fName)
		if err != nil {
			log.Fatalf("File could not be created: %s\n", fName)
		}
		f.WriteString(s)
	}

	// Create workflow
	wf := scipipe.NewWorkflow("wf", 4)

	letterGlobber := NewFileGlobber(wf, "letter_globber", "/tmp/letterfile_*.txt")
	numberGlobber := NewFileGlobber(wf, "number_globber", "/tmp/numberfile_*.txt")

	fileCombiner := NewFileCombinator(wf, "file_combiner")
	fileCombiner.In("letters").From(letterGlobber.Out())
	fileCombiner.In("numbers").From(numberGlobber.Out())

	catenator := wf.NewProc("catenator", "cat {i:letters} {i:numbers} > {o:combined}")
	catenator.In("letters").From(fileCombiner.Out("letters"))
	catenator.In("numbers").From(fileCombiner.Out("numbers"))
	catenator.SetOut("combined", "/tmp/{i:letters|basename|%.txt}.{i:numbers|basename|%.txt}.combined.txt")

	wf.Run()

	for _, l := range []string{"a", "b"} {
		for _, n := range []string{"1", "2", "3"} {
			filePath := fmt.Sprintf("/tmp/letterfile_%s.numberfile_%s.combined.txt", l, n)
			if _, err := os.Stat(filePath); os.IsNotExist(err) {
				log.Fatal("File did not exist: " + filePath)
			}
		}
	}

	// Clean up files
	filePaths := []string{}
	for _, s := range letters {
		filePaths = append(filePaths, fmt.Sprintf("/tmp/letterfile_%s.txt", s))
	}
	for _, s := range numbers {
		filePaths = append(filePaths, fmt.Sprintf("/tmp/numberfile_%s.txt", s))
	}
	for _, l := range []string{"a", "b"} {
		for _, n := range []string{"1", "2", "3"} {
			filePaths = append(filePaths, fmt.Sprintf("/tmp/letterfile_%s.numberfile_%s.combined.txt", l, n))
			filePaths = append(filePaths, filePaths[len(filePaths)-1]+".audit.json")
		}
	}
	for _, filePath := range filePaths {
		err := os.Remove(filePath)
		if err != nil {
			log.Fatal("Could not delete file:", filePath, "\n", err)
		}
	}
}
<commit_msg>Make test of FileCombinator with more than 16 files, to ensure fix of bug with more than 16 files<commit_after>package components

import (
	"fmt"
	"os"
	"testing"

	"log"

	"github.com/scipipe/scipipe"
)

func TestFileCombinator(t *testing.T) {
	letters := []string{"a", "b"}
	numbers := []string{"1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13", "14", "15", "16", "17"}

	// Create letter files
	for _, s := range letters {
		fName := "/tmp/letterfile_" + s + ".txt"
		f, err := os.Create(fName)
		if err != nil {
			log.Fatalf("File could not be created: %s\n", fName)
		}
		f.WriteString(s)
	}

	// Create number files
	for _, s := range numbers {
		fName := "/tmp/numberfile_" + s + ".txt"
		f, err := os.Create(fName)
		if err != nil {
			log.Fatalf("File could not be created: %s\n", fName)
		}
		f.WriteString(s)
	}

	// Create workflow
	wf := scipipe.NewWorkflow("wf", 4)

	letterGlobber := NewFileGlobber(wf, "letter_globber", "/tmp/letterfile_*.txt")
	numberGlobber := NewFileGlobber(wf, "number_globber", "/tmp/numberfile_*.txt")

	fileCombiner := NewFileCombinator(wf, "file_combiner")
	fileCombiner.In("letters").From(letterGlobber.Out())
	fileCombiner.In("numbers").From(numberGlobber.Out())

	catenator := wf.NewProc("catenator", "cat {i:letters} {i:numbers} > {o:combined}")
	catenator.In("letters").From(fileCombiner.Out("letters"))
	catenator.In("numbers").From(fileCombiner.Out("numbers"))
	catenator.SetOut("combined", "/tmp/{i:letters|basename|%.txt}.{i:numbers|basename|%.txt}.combined.txt")

	wf.Run()

	for _, l := range []string{"a", "b"} {
		for _, n := range []string{"1", "2", "3"} {
			filePath := fmt.Sprintf("/tmp/letterfile_%s.numberfile_%s.combined.txt", l, n)
			if _, err := os.Stat(filePath); os.IsNotExist(err) {
				log.Fatal("File did not exist: " + filePath)
			}
		}
	}

	// Clean up files
	filePaths := []string{}
	for _, s := range letters {
		filePaths = append(filePaths, fmt.Sprintf("/tmp/letterfile_%s.txt", s))
	}
	for _, s := range numbers {
		filePaths = append(filePaths, fmt.Sprintf("/tmp/numberfile_%s.txt", s))
	}
	for _, l := range []string{"a", "b"} {
		for _, n := range []string{"1", "2", "3"} {
			filePaths = append(filePaths, fmt.Sprintf("/tmp/letterfile_%s.numberfile_%s.combined.txt", l, n))
			filePaths = append(filePaths, filePaths[len(filePaths)-1]+".audit.json")
		}
	}
	for _, filePath := range filePaths {
		err := os.Remove(filePath)
		if err != nil {
			log.Fatal("Could not delete file:", filePath, "\n", err)
		}
	}
}
<|endoftext|>


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

package cpu

const CacheLinePadSize = 64

// HWCap may be initialized by archauxv and
// should not be changed after it was initialized.
var HWCap uint

// HWCAP bits. These are exposed by Linux.
const (
	hwcap_AES     = 1 << 3
	hwcap_PMULL   = 1 << 4
	hwcap_SHA1    = 1 << 5
	hwcap_SHA2    = 1 << 6
	hwcap_CRC32   = 1 << 7
	hwcap_ATOMICS = 1 << 8
	hwcap_CPUID   = 1 << 11
)

func doinit() {
	options = []option{
		{Name: "aes", Feature: &ARM64.HasAES},
		{Name: "pmull", Feature: &ARM64.HasPMULL},
		{Name: "sha1", Feature: &ARM64.HasSHA1},
		{Name: "sha2", Feature: &ARM64.HasSHA2},
		{Name: "crc32", Feature: &ARM64.HasCRC32},
		{Name: "atomics", Feature: &ARM64.HasATOMICS},
		{Name: "cpuid", Feature: &ARM64.HasCPUID},
		{Name: "isNeoverseN1", Feature: &ARM64.IsNeoverseN1},
		{Name: "isZeus", Feature: &ARM64.IsZeus},
	}

	switch GOOS {
	case "linux", "android":
		// HWCap was populated by the runtime from the auxillary vector.
		// Use HWCap information since reading aarch64 system registers
		// is not supported in user space on older linux kernels.
		ARM64.HasAES = isSet(HWCap, hwcap_AES)
		ARM64.HasPMULL = isSet(HWCap, hwcap_PMULL)
		ARM64.HasSHA1 = isSet(HWCap, hwcap_SHA1)
		ARM64.HasSHA2 = isSet(HWCap, hwcap_SHA2)
		ARM64.HasCRC32 = isSet(HWCap, hwcap_CRC32)
		ARM64.HasCPUID = isSet(HWCap, hwcap_CPUID)

		// The Samsung S9+ kernel reports support for atomics, but not all cores
		// actually support them, resulting in SIGILL. See issue #28431.
		// TODO(elias.naur): Only disable the optimization on bad chipsets on android.
		ARM64.HasATOMICS = isSet(HWCap, hwcap_ATOMICS) && GOOS != "android"

		// Check to see if executing on a NeoverseN1 and in order to do that,
		// check the AUXV for the CPUID bit. The getMIDR function executes an
		// instruction which would normally be an illegal instruction, but it's
		// trapped by the kernel, the value sanitized and then returned. Without
		// the CPUID bit the kernel will not trap the instruction and the process
		// will be terminated with SIGILL.
		if ARM64.HasCPUID {
			midr := getMIDR()
			part_num := uint16((midr >> 4) & 0xfff)
			implementor := byte((midr >> 24) & 0xff)

			if implementor == 'A' && part_num == 0xd0c {
				ARM64.IsNeoverseN1 = true
			}
			if implementor == 'A' && part_num == 0xd40 {
				ARM64.IsZeus = true
			}
		}

	case "freebsd":
		// Retrieve info from system register ID_AA64ISAR0_EL1.
		isar0 := getisar0()

		// ID_AA64ISAR0_EL1
		switch extractBits(isar0, 4, 7) {
		case 1:
			ARM64.HasAES = true
		case 2:
			ARM64.HasAES = true
			ARM64.HasPMULL = true
		}

		switch extractBits(isar0, 8, 11) {
		case 1:
			ARM64.HasSHA1 = true
		}

		switch extractBits(isar0, 12, 15) {
		case 1, 2:
			ARM64.HasSHA2 = true
		}

		switch extractBits(isar0, 16, 19) {
		case 1:
			ARM64.HasCRC32 = true
		}

		switch extractBits(isar0, 20, 23) {
		case 2:
			ARM64.HasATOMICS = true
		}
	default:
		// Other operating systems do not support reading HWCap from auxillary vector
		// or reading privileged aarch64 system registers in user space.
	}
}

func extractBits(data uint64, start, end uint) uint {
	return (uint)(data>>start) & ((1 << (end - start + 1)) - 1)
}

func isSet(hwc uint, value uint) bool {
	return hwc&value != 0
}

func getisar0() uint64

func getMIDR() uint64
<commit_msg>internal/cpu: fix typo in cpu_arm64.go<commit_after>// Copyright 2017 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package cpu

const CacheLinePadSize = 64

// HWCap may be initialized by archauxv and
// should not be changed after it was initialized.
var HWCap uint

// HWCAP bits. These are exposed by Linux.
const (
	hwcap_AES     = 1 << 3
	hwcap_PMULL   = 1 << 4
	hwcap_SHA1    = 1 << 5
	hwcap_SHA2    = 1 << 6
	hwcap_CRC32   = 1 << 7
	hwcap_ATOMICS = 1 << 8
	hwcap_CPUID   = 1 << 11
)

func doinit() {
	options = []option{
		{Name: "aes", Feature: &ARM64.HasAES},
		{Name: "pmull", Feature: &ARM64.HasPMULL},
		{Name: "sha1", Feature: &ARM64.HasSHA1},
		{Name: "sha2", Feature: &ARM64.HasSHA2},
		{Name: "crc32", Feature: &ARM64.HasCRC32},
		{Name: "atomics", Feature: &ARM64.HasATOMICS},
		{Name: "cpuid", Feature: &ARM64.HasCPUID},
		{Name: "isNeoverseN1", Feature: &ARM64.IsNeoverseN1},
		{Name: "isZeus", Feature: &ARM64.IsZeus},
	}

	switch GOOS {
	case "linux", "android":
		// HWCap was populated by the runtime from the auxiliary vector.
		// Use HWCap information since reading aarch64 system registers
		// is not supported in user space on older linux kernels.
		ARM64.HasAES = isSet(HWCap, hwcap_AES)
		ARM64.HasPMULL = isSet(HWCap, hwcap_PMULL)
		ARM64.HasSHA1 = isSet(HWCap, hwcap_SHA1)
		ARM64.HasSHA2 = isSet(HWCap, hwcap_SHA2)
		ARM64.HasCRC32 = isSet(HWCap, hwcap_CRC32)
		ARM64.HasCPUID = isSet(HWCap, hwcap_CPUID)

		// The Samsung S9+ kernel reports support for atomics, but not all cores
		// actually support them, resulting in SIGILL. See issue #28431.
		// TODO(elias.naur): Only disable the optimization on bad chipsets on android.
		ARM64.HasATOMICS = isSet(HWCap, hwcap_ATOMICS) && GOOS != "android"

		// Check to see if executing on a NeoverseN1 and in order to do that,
		// check the AUXV for the CPUID bit. The getMIDR function executes an
		// instruction which would normally be an illegal instruction, but it's
		// trapped by the kernel, the value sanitized and then returned. Without
		// the CPUID bit the kernel will not trap the instruction and the process
		// will be terminated with SIGILL.
		if ARM64.HasCPUID {
			midr := getMIDR()
			part_num := uint16((midr >> 4) & 0xfff)
			implementor := byte((midr >> 24) & 0xff)

			if implementor == 'A' && part_num == 0xd0c {
				ARM64.IsNeoverseN1 = true
			}
			if implementor == 'A' && part_num == 0xd40 {
				ARM64.IsZeus = true
			}
		}

	case "freebsd":
		// Retrieve info from system register ID_AA64ISAR0_EL1.
		isar0 := getisar0()

		// ID_AA64ISAR0_EL1
		switch extractBits(isar0, 4, 7) {
		case 1:
			ARM64.HasAES = true
		case 2:
			ARM64.HasAES = true
			ARM64.HasPMULL = true
		}

		switch extractBits(isar0, 8, 11) {
		case 1:
			ARM64.HasSHA1 = true
		}

		switch extractBits(isar0, 12, 15) {
		case 1, 2:
			ARM64.HasSHA2 = true
		}

		switch extractBits(isar0, 16, 19) {
		case 1:
			ARM64.HasCRC32 = true
		}

		switch extractBits(isar0, 20, 23) {
		case 2:
			ARM64.HasATOMICS = true
		}
	default:
		// Other operating systems do not support reading HWCap from auxiliary vector
		// or reading privileged aarch64 system registers in user space.
	}
}

func extractBits(data uint64, start, end uint) uint {
	return (uint)(data>>start) & ((1 << (end - start + 1)) - 1)
}

func isSet(hwc uint, value uint) bool {
	return hwc&value != 0
}

func getisar0() uint64

func getMIDR() uint64
<|endoftext|>


<commit_before>// Copyright (C) 2017 Jennal(jennalcn@gmail.com). All rights reserved.
//
// Licensed under the MIT License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// http://opensource.org/licenses/MIT
//
// 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 heartbeat

import (
	"fmt"
	"sync"
	"time"

	"github.com/jennal/goplay/consts"
	"github.com/jennal/goplay/log"
	"github.com/jennal/goplay/pkg"
	"github.com/jennal/goplay/session"
)

const (
	TIMEOUT     time.Duration = 3000 * time.Millisecond
	INTERNAL    time.Duration = consts.HeartBeatRate * time.Millisecond
	MAX_TIMEOUT int           = 3
)

type HeartBeatProcessor struct {
	mapMutex     sync.Mutex
	timeoutMutex sync.Mutex

	manager *HeartBeatManager
	sess    *session.Session
	times   map[pkg.PackageIDType]time.Time

	lastPing     time.Duration
	avgPing      time.Duration
	minPing      time.Duration
	maxPing      time.Duration
	sendCount    int64
	recvCount    int64
	timeoutCount int

	stopSignal bool
}

func NewHeartBeatProcessor(manager *HeartBeatManager) *HeartBeatProcessor {
	result := &HeartBeatProcessor{
		manager: manager,
		times:   make(map[pkg.PackageIDType]time.Time),
	}
	go result.checkTimeOut()

	return result
}

func (self *HeartBeatProcessor) Stop() {
	self.stopSignal = true
}

func (self *HeartBeatProcessor) pushTime(id pkg.PackageIDType, t time.Time) {
	self.mapMutex.Lock()
	defer self.mapMutex.Unlock()

	self.times[id] = t
}

func (self *HeartBeatProcessor) popTime(id pkg.PackageIDType) *time.Time {
	self.mapMutex.Lock()
	defer self.mapMutex.Unlock()

	item, ok := self.times[id]
	if !ok {
		return nil
	}
	delete(self.times, id)

	return &item
}

func (self *HeartBeatProcessor) checkTimeOut() {
	for {
		if self.stopSignal {
			return
		}

		// fmt.Printf("[%v] Check Timeout\n", time.Now())
		ids := []pkg.PackageIDType{}
		now := time.Now()

		for id, val := range self.times {
			if now.Sub(val) > TIMEOUT {
				// fmt.Printf("\t => Timeout: %v => %v, id = %v\n", now, val, id)
				ids = append(ids, id)
			}
		}

		for _, id := range ids {
			self.popTime(id)
			if !self.incTimeOut() {
				return
			}
		}

		time.Sleep(TIMEOUT)
	}
}

func (self *HeartBeatProcessor) incTimeOut() bool {
	self.timeoutMutex.Lock()
	defer self.timeoutMutex.Unlock()

	self.timeoutCount++
	// fmt.Println("TimeOut:", self.sess, self.timeoutCount, MAX_TIMEOUT)
	if self.timeoutCount >= MAX_TIMEOUT {
		self.sess.Disconnect()
		return false
	}

	return true
}

func (self *HeartBeatProcessor) resetTimeOut() {
	self.timeoutMutex.Lock()
	defer self.timeoutMutex.Unlock()

	self.timeoutCount = 0
	// fmt.Println("Reset timeOut:", self.sess, self.timeoutCount, MAX_TIMEOUT)
}

func (self *HeartBeatProcessor) newHeader(sess *session.Session) *pkg.Header {
	h := sess.NewHeartBeatHeader()
	self.pushTime(h.ID, time.Now())
	return h
}

func (self *HeartBeatProcessor) SetClient(sess *session.Session) bool /* return false to ignore */ {
	self.sess = sess

	go func() {
		for {
			if self.stopSignal {
				return
			}

			if err := sess.Send(self.newHeader(sess), []byte{}); err == nil {
				self.sendCount++
				self.manager.incSendCount()
			} else {
				sess.Disconnect()
				return
			}

			time.Sleep(INTERNAL)
		}
	}()

	return true
}

func (self *HeartBeatProcessor) OnRecv(sess *session.Session, header *pkg.Header, body []byte) bool /* return false to ignore */ {
	if header.Type != pkg.PKG_HEARTBEAT && header.Type != pkg.PKG_HEARTBEAT_RESPONSE {
		return true
	}

	if sess != self.sess {
		log.Errorf("can't come to here, must be something wrong")
		return true
	}

	self.resetTimeOut()

	switch header.Type {
	case pkg.PKG_HEARTBEAT:
		resp := sess.NewHeartBeatResponseHeader(header)
		sess.Send(resp, []byte{})
	case pkg.PKG_HEARTBEAT_RESPONSE:
		lastTime := self.popTime(header.ID)
		if lastTime == nil {
			return false
		}

		self.recvCount++
		self.manager.incRecvCount()
		self.lastPing = time.Since(*lastTime)
		self.avgPing = (time.Duration)((float64(self.avgPing)*float64(self.recvCount-1) + float64(self.lastPing)) / float64(self.recvCount))
		if self.minPing == 0 || self.lastPing < self.minPing {
			self.minPing = self.lastPing
		}

		if self.maxPing == 0 || self.lastPing > self.maxPing {
			self.maxPing = self.lastPing
		}

		// fmt.Println(self.Info())
	}

	return false
}

func (self *HeartBeatProcessor) Info() string {
	return fmt.Sprintf("Heart Beat For Sess(%d) Statistics ==========\n", self.sess.ID) +
		fmt.Sprintf("=> Last Ping: %v\n", self.lastPing/time.Millisecond) +
		fmt.Sprintf("=> Min Ping: %v\n", self.minPing/time.Millisecond) +
		fmt.Sprintf("=> Max Ping: %v\n", self.maxPing/time.Millisecond) +
		fmt.Sprintf("=> Avg Ping: %v\n", self.avgPing/time.Millisecond) +
		fmt.Sprintf("=> Send Count: %v\n", self.sendCount) +
		fmt.Sprintf("=> Recv Count: %v\n", self.recvCount) +
		fmt.Sprintf("=> Lost Rate: %.2g %%", 100-100*float64(self.recvCount)/float64(self.sendCount))
}
<commit_msg>add lock to heartbeat processor checkTimeOut<commit_after>// Copyright (C) 2017 Jennal(jennalcn@gmail.com). All rights reserved.
//
// Licensed under the MIT License (the "License"); you may not use this file except
// in compliance with the License. You may obtain a copy of the License at
//
// http://opensource.org/licenses/MIT
//
// 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 heartbeat

import (
	"fmt"
	"sync"
	"time"

	"github.com/jennal/goplay/consts"
	"github.com/jennal/goplay/log"
	"github.com/jennal/goplay/pkg"
	"github.com/jennal/goplay/session"
)

const (
	TIMEOUT     time.Duration = 3000 * time.Millisecond
	INTERNAL    time.Duration = consts.HeartBeatRate * time.Millisecond
	MAX_TIMEOUT int           = 3
)

type HeartBeatProcessor struct {
	mapMutex     sync.Mutex
	timeoutMutex sync.Mutex

	manager *HeartBeatManager
	sess    *session.Session
	times   map[pkg.PackageIDType]time.Time

	lastPing     time.Duration
	avgPing      time.Duration
	minPing      time.Duration
	maxPing      time.Duration
	sendCount    int64
	recvCount    int64
	timeoutCount int

	stopSignal bool
}

func NewHeartBeatProcessor(manager *HeartBeatManager) *HeartBeatProcessor {
	result := &HeartBeatProcessor{
		manager: manager,
		times:   make(map[pkg.PackageIDType]time.Time),
	}
	go result.checkTimeOut()

	return result
}

func (self *HeartBeatProcessor) Stop() {
	self.stopSignal = true
}

func (self *HeartBeatProcessor) pushTime(id pkg.PackageIDType, t time.Time) {
	self.mapMutex.Lock()
	defer self.mapMutex.Unlock()

	self.times[id] = t
}

func (self *HeartBeatProcessor) popTime(id pkg.PackageIDType) *time.Time {
	self.mapMutex.Lock()
	defer self.mapMutex.Unlock()

	item, ok := self.times[id]
	if !ok {
		return nil
	}
	delete(self.times, id)

	return &item
}

func (self *HeartBeatProcessor) checkTimeOut() {
	for {
		if self.stopSignal {
			return
		}

		// fmt.Printf("[%v] Check Timeout\n", time.Now())
		ids := []pkg.PackageIDType{}
		now := time.Now()

		self.mapMutex.Lock()
		for id, val := range self.times {
			if now.Sub(val) > TIMEOUT {
				// fmt.Printf("\t => Timeout: %v => %v, id = %v\n", now, val, id)
				ids = append(ids, id)
			}
		}
		self.mapMutex.Unlock()

		for _, id := range ids {
			self.popTime(id)
			if !self.incTimeOut() {
				return
			}
		}

		time.Sleep(TIMEOUT)
	}
}

func (self *HeartBeatProcessor) incTimeOut() bool {
	self.timeoutMutex.Lock()
	defer self.timeoutMutex.Unlock()

	self.timeoutCount++
	// fmt.Println("TimeOut:", self.sess, self.timeoutCount, MAX_TIMEOUT)
	if self.timeoutCount >= MAX_TIMEOUT {
		self.sess.Disconnect()
		return false
	}

	return true
}

func (self *HeartBeatProcessor) resetTimeOut() {
	self.timeoutMutex.Lock()
	defer self.timeoutMutex.Unlock()

	self.timeoutCount = 0
	// fmt.Println("Reset timeOut:", self.sess, self.timeoutCount, MAX_TIMEOUT)
}

func (self *HeartBeatProcessor) newHeader(sess *session.Session) *pkg.Header {
	h := sess.NewHeartBeatHeader()
	self.pushTime(h.ID, time.Now())
	return h
}

func (self *HeartBeatProcessor) SetClient(sess *session.Session) bool /* return false to ignore */ {
	self.sess = sess

	go func() {
		for {
			if self.stopSignal {
				return
			}

			if err := sess.Send(self.newHeader(sess), []byte{}); err == nil {
				self.sendCount++
				self.manager.incSendCount()
			} else {
				sess.Disconnect()
				return
			}

			time.Sleep(INTERNAL)
		}
	}()

	return true
}

func (self *HeartBeatProcessor) OnRecv(sess *session.Session, header *pkg.Header, body []byte) bool /* return false to ignore */ {
	if header.Type != pkg.PKG_HEARTBEAT && header.Type != pkg.PKG_HEARTBEAT_RESPONSE {
		return true
	}

	if sess != self.sess {
		log.Errorf("can't come to here, must be something wrong")
		return true
	}

	self.resetTimeOut()

	switch header.Type {
	case pkg.PKG_HEARTBEAT:
		resp := sess.NewHeartBeatResponseHeader(header)
		sess.Send(resp, []byte{})
	case pkg.PKG_HEARTBEAT_RESPONSE:
		lastTime := self.popTime(header.ID)
		if lastTime == nil {
			return false
		}

		self.recvCount++
		self.manager.incRecvCount()
		self.lastPing = time.Since(*lastTime)
		self.avgPing = (time.Duration)((float64(self.avgPing)*float64(self.recvCount-1) + float64(self.lastPing)) / float64(self.recvCount))
		if self.minPing == 0 || self.lastPing < self.minPing {
			self.minPing = self.lastPing
		}

		if self.maxPing == 0 || self.lastPing > self.maxPing {
			self.maxPing = self.lastPing
		}

		// fmt.Println(self.Info())
	}

	return false
}

func (self *HeartBeatProcessor) Info() string {
	return fmt.Sprintf("Heart Beat For Sess(%d) Statistics ==========\n", self.sess.ID) +
		fmt.Sprintf("=> Last Ping: %v\n", self.lastPing/time.Millisecond) +
		fmt.Sprintf("=> Min Ping: %v\n", self.minPing/time.Millisecond) +
		fmt.Sprintf("=> Max Ping: %v\n", self.maxPing/time.Millisecond) +
		fmt.Sprintf("=> Avg Ping: %v\n", self.avgPing/time.Millisecond) +
		fmt.Sprintf("=> Send Count: %v\n", self.sendCount) +
		fmt.Sprintf("=> Recv Count: %v\n", self.recvCount) +
		fmt.Sprintf("=> Lost Rate: %.2g %%", 100-100*float64(self.recvCount)/float64(self.sendCount))
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package framework

import (
	"context"
	"flag"
	"net"
	"net/http"
	"net/http/httptest"
	"path"
	"strconv"
	"time"

	"github.com/google/uuid"

	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/util/wait"
	authauthenticator "k8s.io/apiserver/pkg/authentication/authenticator"
	"k8s.io/apiserver/pkg/authentication/authenticatorfactory"
	authenticatorunion "k8s.io/apiserver/pkg/authentication/request/union"
	"k8s.io/apiserver/pkg/authentication/user"
	"k8s.io/apiserver/pkg/authorization/authorizer"
	"k8s.io/apiserver/pkg/authorization/authorizerfactory"
	authorizerunion "k8s.io/apiserver/pkg/authorization/union"
	openapinamer "k8s.io/apiserver/pkg/endpoints/openapi"
	genericfeatures "k8s.io/apiserver/pkg/features"
	genericapiserver "k8s.io/apiserver/pkg/server"
	"k8s.io/apiserver/pkg/server/options"
	serverstorage "k8s.io/apiserver/pkg/server/storage"
	"k8s.io/apiserver/pkg/storage/storagebackend"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
	utilflowcontrol "k8s.io/apiserver/pkg/util/flowcontrol"
	"k8s.io/client-go/informers"
	clientset "k8s.io/client-go/kubernetes"
	restclient "k8s.io/client-go/rest"
	"k8s.io/component-base/version"
	"k8s.io/klog/v2"
	openapicommon "k8s.io/kube-openapi/pkg/common"
	"k8s.io/kube-openapi/pkg/validation/spec"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/controlplane"
	"k8s.io/kubernetes/pkg/generated/openapi"
	"k8s.io/kubernetes/pkg/kubeapiserver"
	kubeletclient "k8s.io/kubernetes/pkg/kubelet/client"
)

// Config is a struct of configuration directives for NewControlPlaneComponents.
type Config struct {
	// If nil, a default is used, partially filled configs will not get populated.
	InstanceConfig          *controlplane.Config
	StartReplicationManager bool
	// Client throttling qps
	QPS float32
	// Client burst qps, also burst replicas allowed in rc manager
	Burst int
	// TODO: Add configs for endpoints controller, scheduler etc
}

// alwaysAllow always allows an action
type alwaysAllow struct{}

func (alwaysAllow) Authorize(ctx context.Context, requestAttributes authorizer.Attributes) (authorizer.Decision, string, error) {
	return authorizer.DecisionAllow, "always allow", nil
}

// alwaysEmpty simulates "no authentication" for old tests
func alwaysEmpty(req *http.Request) (*authauthenticator.Response, bool, error) {
	return &authauthenticator.Response{
		User: &user.DefaultInfo{
			Name: "",
		},
	}, true, nil
}

// APIServerReceiver can be used to provide the API server to a custom incoming server function
type APIServerReceiver interface {
	SetAPIServer(m *controlplane.Instance)
}

// APIServerHolder implements
type APIServerHolder struct {
	Initialized chan struct{}
	M           *controlplane.Instance
}

// SetAPIServer assigns the current API server.
func (h *APIServerHolder) SetAPIServer(m *controlplane.Instance) {
	h.M = m
	close(h.Initialized)
}

// DefaultOpenAPIConfig returns an openapicommon.Config initialized to default values.
func DefaultOpenAPIConfig() *openapicommon.Config {
	openAPIConfig := genericapiserver.DefaultOpenAPIConfig(openapi.GetOpenAPIDefinitions, openapinamer.NewDefinitionNamer(legacyscheme.Scheme))
	openAPIConfig.Info = &spec.Info{
		InfoProps: spec.InfoProps{
			Title:   "Kubernetes",
			Version: "unversioned",
		},
	}
	openAPIConfig.DefaultResponse = &spec.Response{
		ResponseProps: spec.ResponseProps{
			Description: "Default Response.",
		},
	}
	openAPIConfig.GetDefinitions = openapi.GetOpenAPIDefinitions

	return openAPIConfig
}

// startAPIServerOrDie starts a kubernetes API server and an httpserver to handle api requests
func startAPIServerOrDie(controlPlaneConfig *controlplane.Config, incomingServer *httptest.Server, apiServerReceiver APIServerReceiver) (*controlplane.Instance, *httptest.Server, CloseFunc) {
	var m *controlplane.Instance
	var s *httptest.Server

	// Ensure we log at least level 4
	v := flag.Lookup("v").Value
	level, _ := strconv.Atoi(v.String())
	if level < 4 {
		v.Set("4")
	}

	if incomingServer != nil {
		s = incomingServer
	} else {
		s = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			m.GenericAPIServer.Handler.ServeHTTP(w, req)
		}))
	}

	stopCh := make(chan struct{})
	closeFn := func() {
		if m != nil {
			m.GenericAPIServer.RunPreShutdownHooks()
		}
		close(stopCh)
		s.Close()
	}

	if controlPlaneConfig == nil {
		controlPlaneConfig = NewControlPlaneConfig()
		controlPlaneConfig.GenericConfig.OpenAPIConfig = DefaultOpenAPIConfig()
	}

	// set the loopback client config
	if controlPlaneConfig.GenericConfig.LoopbackClientConfig == nil {
		controlPlaneConfig.GenericConfig.LoopbackClientConfig = &restclient.Config{QPS: 50, Burst: 100, ContentConfig: restclient.ContentConfig{NegotiatedSerializer: legacyscheme.Codecs}}
	}
	controlPlaneConfig.GenericConfig.LoopbackClientConfig.Host = s.URL

	privilegedLoopbackToken := uuid.New().String()
	// wrap any available authorizer
	tokens := make(map[string]*user.DefaultInfo)
	tokens[privilegedLoopbackToken] = &user.DefaultInfo{
		Name:   user.APIServerUser,
		UID:    uuid.New().String(),
		Groups: []string{user.SystemPrivilegedGroup},
	}

	tokenAuthenticator := authenticatorfactory.NewFromTokens(tokens, controlPlaneConfig.GenericConfig.Authentication.APIAudiences)
	if controlPlaneConfig.GenericConfig.Authentication.Authenticator == nil {
		controlPlaneConfig.GenericConfig.Authentication.Authenticator = authenticatorunion.New(tokenAuthenticator, authauthenticator.RequestFunc(alwaysEmpty))
	} else {
		controlPlaneConfig.GenericConfig.Authentication.Authenticator = authenticatorunion.New(tokenAuthenticator, controlPlaneConfig.GenericConfig.Authentication.Authenticator)
	}

	if controlPlaneConfig.GenericConfig.Authorization.Authorizer != nil {
		tokenAuthorizer := authorizerfactory.NewPrivilegedGroups(user.SystemPrivilegedGroup)
		controlPlaneConfig.GenericConfig.Authorization.Authorizer = authorizerunion.New(tokenAuthorizer, controlPlaneConfig.GenericConfig.Authorization.Authorizer)
	} else {
		controlPlaneConfig.GenericConfig.Authorization.Authorizer = alwaysAllow{}
	}

	controlPlaneConfig.GenericConfig.LoopbackClientConfig.BearerToken = privilegedLoopbackToken

	clientset, err := clientset.NewForConfig(controlPlaneConfig.GenericConfig.LoopbackClientConfig)
	if err != nil {
		klog.Fatal(err)
	}

	controlPlaneConfig.ExtraConfig.VersionedInformers = informers.NewSharedInformerFactory(clientset, controlPlaneConfig.GenericConfig.LoopbackClientConfig.Timeout)

	if utilfeature.DefaultFeatureGate.Enabled(genericfeatures.APIPriorityAndFairness) {
		controlPlaneConfig.GenericConfig.FlowControl = utilflowcontrol.New(
			controlPlaneConfig.ExtraConfig.VersionedInformers,
			clientset.FlowcontrolV1beta1(),
			controlPlaneConfig.GenericConfig.MaxRequestsInFlight+controlPlaneConfig.GenericConfig.MaxMutatingRequestsInFlight,
			controlPlaneConfig.GenericConfig.RequestTimeout/4,
		)
	}

	if controlPlaneConfig.ExtraConfig.ServiceIPRange.IP == nil {
		controlPlaneConfig.ExtraConfig.ServiceIPRange = net.IPNet{IP: net.ParseIP("10.0.0.0"), Mask: net.CIDRMask(24, 32)}
	}
	m, err = controlPlaneConfig.Complete().New(genericapiserver.NewEmptyDelegate())
	if err != nil {
		// We log the error first so that even if closeFn crashes, the error is shown
		klog.Errorf("error in bringing up the apiserver: %v", err)
		closeFn()
		klog.Fatalf("error in bringing up the apiserver: %v", err)
	}
	if apiServerReceiver != nil {
		apiServerReceiver.SetAPIServer(m)
	}

	// TODO have this start method actually use the normal start sequence for the API server
	// this method never actually calls the `Run` method for the API server
	// fire the post hooks ourselves
	m.GenericAPIServer.PrepareRun()
	m.GenericAPIServer.RunPostStartHooks(stopCh)

	cfg := *controlPlaneConfig.GenericConfig.LoopbackClientConfig
	cfg.ContentConfig.GroupVersion = &schema.GroupVersion{}
	privilegedClient, err := restclient.RESTClientFor(&cfg)
	if err != nil {
		closeFn()
		klog.Fatal(err)
	}
	var lastHealthContent []byte
	err = wait.PollImmediate(100*time.Millisecond, 30*time.Second, func() (bool, error) {
		result := privilegedClient.Get().AbsPath("/healthz").Do(context.TODO())
		status := 0
		result.StatusCode(&status)
		if status == 200 {
			return true, nil
		}
		lastHealthContent, _ = result.Raw()
		return false, nil
	})
	if err != nil {
		closeFn()
		klog.Errorf("last health content: %q", string(lastHealthContent))
		klog.Fatal(err)
	}

	return m, s, closeFn
}

// NewIntegrationTestControlPlaneConfig returns the control plane config appropriate for most integration tests.
func NewIntegrationTestControlPlaneConfig() *controlplane.Config {
	return NewIntegrationTestControlPlaneConfigWithOptions(&ControlPlaneConfigOptions{})
}

// NewIntegrationTestControlPlaneConfigWithOptions returns the control plane config appropriate for most integration tests
// configured with the provided options.
func NewIntegrationTestControlPlaneConfigWithOptions(opts *ControlPlaneConfigOptions) *controlplane.Config {
	controlPlaneConfig := NewControlPlaneConfigWithOptions(opts)
	controlPlaneConfig.GenericConfig.PublicAddress = net.ParseIP("192.168.10.4")
	controlPlaneConfig.ExtraConfig.APIResourceConfigSource = controlplane.DefaultAPIResourceConfigSource()

	// TODO: get rid of these tests or port them to secure serving
	controlPlaneConfig.GenericConfig.SecureServing = &genericapiserver.SecureServingInfo{Listener: fakeLocalhost443Listener{}}

	return controlPlaneConfig
}

// ControlPlaneConfigOptions are the configurable options for a new integration test control plane config.
type ControlPlaneConfigOptions struct {
	EtcdOptions *options.EtcdOptions
}

// DefaultEtcdOptions are the default EtcdOptions for use with integration tests.
func DefaultEtcdOptions() *options.EtcdOptions {
	// This causes the integration tests to exercise the etcd
	// prefix code, so please don't change without ensuring
	// sufficient coverage in other ways.
	etcdOptions := options.NewEtcdOptions(storagebackend.NewDefaultConfig(uuid.New().String(), nil))
	etcdOptions.StorageConfig.Transport.ServerList = []string{GetEtcdURL()}
	return etcdOptions
}

// NewControlPlaneConfig returns a basic control plane config.
func NewControlPlaneConfig() *controlplane.Config {
	return NewControlPlaneConfigWithOptions(&ControlPlaneConfigOptions{})
}

// NewControlPlaneConfigWithOptions returns a basic control plane config configured with the provided options.
func NewControlPlaneConfigWithOptions(opts *ControlPlaneConfigOptions) *controlplane.Config {
	etcdOptions := DefaultEtcdOptions()
	if opts.EtcdOptions != nil {
		etcdOptions = opts.EtcdOptions
	}

	storageConfig := kubeapiserver.NewStorageFactoryConfig()
	storageConfig.APIResourceConfig = serverstorage.NewResourceConfig()
	completedStorageConfig, err := storageConfig.Complete(etcdOptions)
	if err != nil {
		panic(err)
	}
	storageFactory, err := completedStorageConfig.New()
	if err != nil {
		panic(err)
	}

	genericConfig := genericapiserver.NewConfig(legacyscheme.Codecs)
	kubeVersion := version.Get()
	if len(kubeVersion.Major) == 0 {
		kubeVersion.Major = "1"
	}
	if len(kubeVersion.Minor) == 0 {
		kubeVersion.Minor = "22"
	}
	genericConfig.Version = &kubeVersion
	genericConfig.Authorization.Authorizer = authorizerfactory.NewAlwaysAllowAuthorizer()

	// TODO: get rid of these tests or port them to secure serving
	genericConfig.SecureServing = &genericapiserver.SecureServingInfo{Listener: fakeLocalhost443Listener{}}

	err = etcdOptions.ApplyWithStorageFactoryTo(storageFactory, genericConfig)
	if err != nil {
		panic(err)
	}

	return &controlplane.Config{
		GenericConfig: genericConfig,
		ExtraConfig: controlplane.ExtraConfig{
			APIResourceConfigSource: controlplane.DefaultAPIResourceConfigSource(),
			StorageFactory:          storageFactory,
			KubeletClientConfig:     kubeletclient.KubeletClientConfig{Port: 10250},
			APIServerServicePort:    443,
			MasterCount:             1,
		},
	}
}

// CloseFunc can be called to cleanup the API server
type CloseFunc func()

// RunAnAPIServer starts a API server with the provided config.
func RunAnAPIServer(controlPlaneConfig *controlplane.Config) (*controlplane.Instance, *httptest.Server, CloseFunc) {
	if controlPlaneConfig == nil {
		controlPlaneConfig = NewControlPlaneConfig()
		controlPlaneConfig.GenericConfig.EnableProfiling = true
	}
	return startAPIServerOrDie(controlPlaneConfig, nil, nil)
}

// RunAnAPIServerUsingServer starts up an instance using the provided config on the specified server.
func RunAnAPIServerUsingServer(controlPlaneConfig *controlplane.Config, s *httptest.Server, apiServerReceiver APIServerReceiver) (*controlplane.Instance, *httptest.Server, CloseFunc) {
	return startAPIServerOrDie(controlPlaneConfig, s, apiServerReceiver)
}

// SharedEtcd creates a storage config for a shared etcd instance, with a unique prefix.
func SharedEtcd() *storagebackend.Config {
	cfg := storagebackend.NewDefaultConfig(path.Join(uuid.New().String(), "registry"), nil)
	cfg.Transport.ServerList = []string{GetEtcdURL()}
	return cfg
}

type fakeLocalhost443Listener struct{}

func (fakeLocalhost443Listener) Accept() (net.Conn, error) {
	return nil, nil
}

func (fakeLocalhost443Listener) Close() error {
	return nil
}

func (fakeLocalhost443Listener) Addr() net.Addr {
	return &net.TCPAddr{
		IP:   net.IPv4(127, 0, 0, 1),
		Port: 443,
	}
}
<commit_msg>integration test: provide a timeout for /health<commit_after>/*
Copyright 2014 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package framework

import (
	"context"
	"flag"
	"net"
	"net/http"
	"net/http/httptest"
	"path"
	"strconv"
	"time"

	"github.com/google/uuid"

	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/util/wait"
	authauthenticator "k8s.io/apiserver/pkg/authentication/authenticator"
	"k8s.io/apiserver/pkg/authentication/authenticatorfactory"
	authenticatorunion "k8s.io/apiserver/pkg/authentication/request/union"
	"k8s.io/apiserver/pkg/authentication/user"
	"k8s.io/apiserver/pkg/authorization/authorizer"
	"k8s.io/apiserver/pkg/authorization/authorizerfactory"
	authorizerunion "k8s.io/apiserver/pkg/authorization/union"
	openapinamer "k8s.io/apiserver/pkg/endpoints/openapi"
	genericfeatures "k8s.io/apiserver/pkg/features"
	genericapiserver "k8s.io/apiserver/pkg/server"
	"k8s.io/apiserver/pkg/server/options"
	serverstorage "k8s.io/apiserver/pkg/server/storage"
	"k8s.io/apiserver/pkg/storage/storagebackend"
	utilfeature "k8s.io/apiserver/pkg/util/feature"
	utilflowcontrol "k8s.io/apiserver/pkg/util/flowcontrol"
	"k8s.io/client-go/informers"
	clientset "k8s.io/client-go/kubernetes"
	restclient "k8s.io/client-go/rest"
	"k8s.io/component-base/version"
	"k8s.io/klog/v2"
	openapicommon "k8s.io/kube-openapi/pkg/common"
	"k8s.io/kube-openapi/pkg/validation/spec"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/controlplane"
	"k8s.io/kubernetes/pkg/generated/openapi"
	"k8s.io/kubernetes/pkg/kubeapiserver"
	kubeletclient "k8s.io/kubernetes/pkg/kubelet/client"
)

// Config is a struct of configuration directives for NewControlPlaneComponents.
type Config struct {
	// If nil, a default is used, partially filled configs will not get populated.
	InstanceConfig          *controlplane.Config
	StartReplicationManager bool
	// Client throttling qps
	QPS float32
	// Client burst qps, also burst replicas allowed in rc manager
	Burst int
	// TODO: Add configs for endpoints controller, scheduler etc
}

// alwaysAllow always allows an action
type alwaysAllow struct{}

func (alwaysAllow) Authorize(ctx context.Context, requestAttributes authorizer.Attributes) (authorizer.Decision, string, error) {
	return authorizer.DecisionAllow, "always allow", nil
}

// alwaysEmpty simulates "no authentication" for old tests
func alwaysEmpty(req *http.Request) (*authauthenticator.Response, bool, error) {
	return &authauthenticator.Response{
		User: &user.DefaultInfo{
			Name: "",
		},
	}, true, nil
}

// APIServerReceiver can be used to provide the API server to a custom incoming server function
type APIServerReceiver interface {
	SetAPIServer(m *controlplane.Instance)
}

// APIServerHolder implements
type APIServerHolder struct {
	Initialized chan struct{}
	M           *controlplane.Instance
}

// SetAPIServer assigns the current API server.
func (h *APIServerHolder) SetAPIServer(m *controlplane.Instance) {
	h.M = m
	close(h.Initialized)
}

// DefaultOpenAPIConfig returns an openapicommon.Config initialized to default values.
func DefaultOpenAPIConfig() *openapicommon.Config {
	openAPIConfig := genericapiserver.DefaultOpenAPIConfig(openapi.GetOpenAPIDefinitions, openapinamer.NewDefinitionNamer(legacyscheme.Scheme))
	openAPIConfig.Info = &spec.Info{
		InfoProps: spec.InfoProps{
			Title:   "Kubernetes",
			Version: "unversioned",
		},
	}
	openAPIConfig.DefaultResponse = &spec.Response{
		ResponseProps: spec.ResponseProps{
			Description: "Default Response.",
		},
	}
	openAPIConfig.GetDefinitions = openapi.GetOpenAPIDefinitions

	return openAPIConfig
}

// startAPIServerOrDie starts a kubernetes API server and an httpserver to handle api requests
func startAPIServerOrDie(controlPlaneConfig *controlplane.Config, incomingServer *httptest.Server, apiServerReceiver APIServerReceiver) (*controlplane.Instance, *httptest.Server, CloseFunc) {
	var m *controlplane.Instance
	var s *httptest.Server

	// Ensure we log at least level 4
	v := flag.Lookup("v").Value
	level, _ := strconv.Atoi(v.String())
	if level < 4 {
		v.Set("4")
	}

	if incomingServer != nil {
		s = incomingServer
	} else {
		s = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
			m.GenericAPIServer.Handler.ServeHTTP(w, req)
		}))
	}

	stopCh := make(chan struct{})
	closeFn := func() {
		if m != nil {
			m.GenericAPIServer.RunPreShutdownHooks()
		}
		close(stopCh)
		s.Close()
	}

	if controlPlaneConfig == nil {
		controlPlaneConfig = NewControlPlaneConfig()
		controlPlaneConfig.GenericConfig.OpenAPIConfig = DefaultOpenAPIConfig()
	}

	// set the loopback client config
	if controlPlaneConfig.GenericConfig.LoopbackClientConfig == nil {
		controlPlaneConfig.GenericConfig.LoopbackClientConfig = &restclient.Config{QPS: 50, Burst: 100, ContentConfig: restclient.ContentConfig{NegotiatedSerializer: legacyscheme.Codecs}}
	}
	controlPlaneConfig.GenericConfig.LoopbackClientConfig.Host = s.URL

	privilegedLoopbackToken := uuid.New().String()
	// wrap any available authorizer
	tokens := make(map[string]*user.DefaultInfo)
	tokens[privilegedLoopbackToken] = &user.DefaultInfo{
		Name:   user.APIServerUser,
		UID:    uuid.New().String(),
		Groups: []string{user.SystemPrivilegedGroup},
	}

	tokenAuthenticator := authenticatorfactory.NewFromTokens(tokens, controlPlaneConfig.GenericConfig.Authentication.APIAudiences)
	if controlPlaneConfig.GenericConfig.Authentication.Authenticator == nil {
		controlPlaneConfig.GenericConfig.Authentication.Authenticator = authenticatorunion.New(tokenAuthenticator, authauthenticator.RequestFunc(alwaysEmpty))
	} else {
		controlPlaneConfig.GenericConfig.Authentication.Authenticator = authenticatorunion.New(tokenAuthenticator, controlPlaneConfig.GenericConfig.Authentication.Authenticator)
	}

	if controlPlaneConfig.GenericConfig.Authorization.Authorizer != nil {
		tokenAuthorizer := authorizerfactory.NewPrivilegedGroups(user.SystemPrivilegedGroup)
		controlPlaneConfig.GenericConfig.Authorization.Authorizer = authorizerunion.New(tokenAuthorizer, controlPlaneConfig.GenericConfig.Authorization.Authorizer)
	} else {
		controlPlaneConfig.GenericConfig.Authorization.Authorizer = alwaysAllow{}
	}

	controlPlaneConfig.GenericConfig.LoopbackClientConfig.BearerToken = privilegedLoopbackToken

	clientset, err := clientset.NewForConfig(controlPlaneConfig.GenericConfig.LoopbackClientConfig)
	if err != nil {
		klog.Fatal(err)
	}

	controlPlaneConfig.ExtraConfig.VersionedInformers = informers.NewSharedInformerFactory(clientset, controlPlaneConfig.GenericConfig.LoopbackClientConfig.Timeout)

	if utilfeature.DefaultFeatureGate.Enabled(genericfeatures.APIPriorityAndFairness) {
		controlPlaneConfig.GenericConfig.FlowControl = utilflowcontrol.New(
			controlPlaneConfig.ExtraConfig.VersionedInformers,
			clientset.FlowcontrolV1beta1(),
			controlPlaneConfig.GenericConfig.MaxRequestsInFlight+controlPlaneConfig.GenericConfig.MaxMutatingRequestsInFlight,
			controlPlaneConfig.GenericConfig.RequestTimeout/4,
		)
	}

	if controlPlaneConfig.ExtraConfig.ServiceIPRange.IP == nil {
		controlPlaneConfig.ExtraConfig.ServiceIPRange = net.IPNet{IP: net.ParseIP("10.0.0.0"), Mask: net.CIDRMask(24, 32)}
	}
	m, err = controlPlaneConfig.Complete().New(genericapiserver.NewEmptyDelegate())
	if err != nil {
		// We log the error first so that even if closeFn crashes, the error is shown
		klog.Errorf("error in bringing up the apiserver: %v", err)
		closeFn()
		klog.Fatalf("error in bringing up the apiserver: %v", err)
	}
	if apiServerReceiver != nil {
		apiServerReceiver.SetAPIServer(m)
	}

	// TODO have this start method actually use the normal start sequence for the API server
	// this method never actually calls the `Run` method for the API server
	// fire the post hooks ourselves
	m.GenericAPIServer.PrepareRun()
	m.GenericAPIServer.RunPostStartHooks(stopCh)

	cfg := *controlPlaneConfig.GenericConfig.LoopbackClientConfig
	cfg.ContentConfig.GroupVersion = &schema.GroupVersion{}
	privilegedClient, err := restclient.RESTClientFor(&cfg)
	if err != nil {
		closeFn()
		klog.Fatal(err)
	}
	var lastHealthContent []byte
	err = wait.PollImmediate(100*time.Millisecond, 30*time.Second, func() (bool, error) {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()

		result := privilegedClient.Get().AbsPath("/healthz").Do(ctx)
		status := 0
		result.StatusCode(&status)
		if status == 200 {
			return true, nil
		}
		lastHealthContent, _ = result.Raw()
		return false, nil
	})
	if err != nil {
		closeFn()
		klog.Errorf("last health content: %q", string(lastHealthContent))
		klog.Fatal(err)
	}

	return m, s, closeFn
}

// NewIntegrationTestControlPlaneConfig returns the control plane config appropriate for most integration tests.
func NewIntegrationTestControlPlaneConfig() *controlplane.Config {
	return NewIntegrationTestControlPlaneConfigWithOptions(&ControlPlaneConfigOptions{})
}

// NewIntegrationTestControlPlaneConfigWithOptions returns the control plane config appropriate for most integration tests
// configured with the provided options.
func NewIntegrationTestControlPlaneConfigWithOptions(opts *ControlPlaneConfigOptions) *controlplane.Config {
	controlPlaneConfig := NewControlPlaneConfigWithOptions(opts)
	controlPlaneConfig.GenericConfig.PublicAddress = net.ParseIP("192.168.10.4")
	controlPlaneConfig.ExtraConfig.APIResourceConfigSource = controlplane.DefaultAPIResourceConfigSource()

	// TODO: get rid of these tests or port them to secure serving
	controlPlaneConfig.GenericConfig.SecureServing = &genericapiserver.SecureServingInfo{Listener: fakeLocalhost443Listener{}}

	return controlPlaneConfig
}

// ControlPlaneConfigOptions are the configurable options for a new integration test control plane config.
type ControlPlaneConfigOptions struct {
	EtcdOptions *options.EtcdOptions
}

// DefaultEtcdOptions are the default EtcdOptions for use with integration tests.
func DefaultEtcdOptions() *options.EtcdOptions {
	// This causes the integration tests to exercise the etcd
	// prefix code, so please don't change without ensuring
	// sufficient coverage in other ways.
	etcdOptions := options.NewEtcdOptions(storagebackend.NewDefaultConfig(uuid.New().String(), nil))
	etcdOptions.StorageConfig.Transport.ServerList = []string{GetEtcdURL()}
	return etcdOptions
}

// NewControlPlaneConfig returns a basic control plane config.
func NewControlPlaneConfig() *controlplane.Config {
	return NewControlPlaneConfigWithOptions(&ControlPlaneConfigOptions{})
}

// NewControlPlaneConfigWithOptions returns a basic control plane config configured with the provided options.
func NewControlPlaneConfigWithOptions(opts *ControlPlaneConfigOptions) *controlplane.Config {
	etcdOptions := DefaultEtcdOptions()
	if opts.EtcdOptions != nil {
		etcdOptions = opts.EtcdOptions
	}

	storageConfig := kubeapiserver.NewStorageFactoryConfig()
	storageConfig.APIResourceConfig = serverstorage.NewResourceConfig()
	completedStorageConfig, err := storageConfig.Complete(etcdOptions)
	if err != nil {
		panic(err)
	}
	storageFactory, err := completedStorageConfig.New()
	if err != nil {
		panic(err)
	}

	genericConfig := genericapiserver.NewConfig(legacyscheme.Codecs)
	kubeVersion := version.Get()
	if len(kubeVersion.Major) == 0 {
		kubeVersion.Major = "1"
	}
	if len(kubeVersion.Minor) == 0 {
		kubeVersion.Minor = "22"
	}
	genericConfig.Version = &kubeVersion
	genericConfig.Authorization.Authorizer = authorizerfactory.NewAlwaysAllowAuthorizer()

	// TODO: get rid of these tests or port them to secure serving
	genericConfig.SecureServing = &genericapiserver.SecureServingInfo{Listener: fakeLocalhost443Listener{}}

	err = etcdOptions.ApplyWithStorageFactoryTo(storageFactory, genericConfig)
	if err != nil {
		panic(err)
	}

	return &controlplane.Config{
		GenericConfig: genericConfig,
		ExtraConfig: controlplane.ExtraConfig{
			APIResourceConfigSource: controlplane.DefaultAPIResourceConfigSource(),
			StorageFactory:          storageFactory,
			KubeletClientConfig:     kubeletclient.KubeletClientConfig{Port: 10250},
			APIServerServicePort:    443,
			MasterCount:             1,
		},
	}
}

// CloseFunc can be called to cleanup the API server
type CloseFunc func()

// RunAnAPIServer starts a API server with the provided config.
func RunAnAPIServer(controlPlaneConfig *controlplane.Config) (*controlplane.Instance, *httptest.Server, CloseFunc) {
	if controlPlaneConfig == nil {
		controlPlaneConfig = NewControlPlaneConfig()
		controlPlaneConfig.GenericConfig.EnableProfiling = true
	}
	return startAPIServerOrDie(controlPlaneConfig, nil, nil)
}

// RunAnAPIServerUsingServer starts up an instance using the provided config on the specified server.
func RunAnAPIServerUsingServer(controlPlaneConfig *controlplane.Config, s *httptest.Server, apiServerReceiver APIServerReceiver) (*controlplane.Instance, *httptest.Server, CloseFunc) {
	return startAPIServerOrDie(controlPlaneConfig, s, apiServerReceiver)
}

// SharedEtcd creates a storage config for a shared etcd instance, with a unique prefix.
func SharedEtcd() *storagebackend.Config {
	cfg := storagebackend.NewDefaultConfig(path.Join(uuid.New().String(), "registry"), nil)
	cfg.Transport.ServerList = []string{GetEtcdURL()}
	return cfg
}

type fakeLocalhost443Listener struct{}

func (fakeLocalhost443Listener) Accept() (net.Conn, error) {
	return nil, nil
}

func (fakeLocalhost443Listener) Close() error {
	return nil
}

func (fakeLocalhost443Listener) Addr() net.Addr {
	return &net.TCPAddr{
		IP:   net.IPv4(127, 0, 0, 1),
		Port: 443,
	}
}
<|endoftext|>


<commit_before>package integrationtests

import (
	"testing"

	"github.com/bunsenapp/go-selenium/src"
)

func Test_ElementAttribute_CanRetrieveAttributeCorrectly(t *testing.T) {
	setUp()
	defer tearDown()

	tests := []struct {
		url  string
		by   goselenium.By
		attr string
		val  string
	}{
		{
			url:  "https://www.google.com",
			by:   goselenium.ByCSSSelector("input#lst-ib"),
			attr: "maxlength",
			val:  "2048",
		},
	}
	for _, te := range tests {
		driver := createDriver(t)
		_, err := driver.CreateSession()
		if err != nil {
			errorAndWrap(t, "Error thrown whilst creating session.", err)
		}

		_, err = driver.Go(te.url)
		if err != nil {
			errorAndWrap(t, "Error thrown whilst visiting url.", err)
		}

		el, err := driver.FindElement(te.by)
		if err != nil || el == nil {
			errorAndWrap(t, "Error whilst finding element or element was not found", err)
		}

		att, err := el.Attribute(te.attr)
		if err != nil || att.Value != te.val {
			errorAndWrap(t, "Error whilst retrieving attribute or attribute value was not correct", err)
		}

		driver.DeleteSession()

		printObjectResult(att)
	}
}
<commit_msg>Fix attribute errors.<commit_after>package integrationtests

import (
	"strings"
	"testing"

	"github.com/bunsenapp/go-selenium/src"
)

func Test_ElementAttribute_CanRetrieveAttributeCorrectly(t *testing.T) {
	setUp()
	defer tearDown()

	tests := []struct {
		url  string
		by   goselenium.By
		attr string
		val  string
	}{
		{
			url:  "https://www.google.com",
			by:   goselenium.ByCSSSelector("input#lst-ib"),
			attr: "maxlength",
			val:  "2048",
		},
		{
			url:  "https://news.ycombinator.com",
			by:   goselenium.ByCSSSelector("b.hnname > a"),
			attr: "href",
			val:  "news",
		},
	}
	for _, te := range tests {
		driver := createDriver(t)
		_, err := driver.CreateSession()
		if err != nil {
			errorAndWrap(t, "Error thrown whilst creating session.", err)
		}

		_, err = driver.Go(te.url)
		if err != nil {
			errorAndWrap(t, "Error thrown whilst visiting url.", err)
		}

		el, err := driver.FindElement(te.by)
		if err != nil || el == nil {
			errorAndWrap(t, "Error whilst finding element or element was not found", err)
		}

		att, err := el.Attribute(te.attr)
		if err != nil || !strings.Contains(att.Value, te.val) {
			errorAndWrap(t, "Error whilst retrieving attribute or attribute value was not correct", err)
		}

		driver.DeleteSession()

		printObjectResult(att)
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 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 writing, software distributed
// under the License is distributed on a "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 contextimpl

import (
	"errors"
	"reflect"
	"sync"
	"time"
)

type Context interface {
	Deadline() (deadline time.Time, ok bool)
	Done() <-chan struct{}
	Err() error
	Value(key interface{}) interface{}
}

type emptyCtx int

func (emptyCtx) Deadline() (deadline time.Time, ok bool) { return }
func (emptyCtx) Done() <-chan struct{}                   { return nil }
func (emptyCtx) Err() error                              { return nil }
func (emptyCtx) Value(key interface{}) interface{}       { return nil }

var (
	background = new(emptyCtx)
	todo       = new(emptyCtx)
)

func Background() Context { return background }
func TODO() Context       { return todo }

type cancelCtx struct {
	Context
	done chan struct{}
	err  error
	mu   sync.Mutex
}

func (ctx *cancelCtx) Done() <-chan struct{} { return ctx.done }
func (ctx *cancelCtx) Err() error {
	ctx.mu.Lock()
	defer ctx.mu.Unlock()
	return ctx.err
}

var Canceled = errors.New("context canceled")

type CancelFunc func()

func WithCancel(parent Context) (Context, CancelFunc) {
	ctx := &cancelCtx{
		Context: parent,
		done:    make(chan struct{}),
	}

	cancel := func() { ctx.cancel(Canceled) }

	go func() {
		select {
		case <-parent.Done():
			ctx.cancel(parent.Err())
		case <-ctx.Done():
		}
	}()

	return ctx, cancel
}

func (ctx *cancelCtx) cancel(err error) {
	ctx.mu.Lock()
	defer ctx.mu.Unlock()
	if ctx.err != nil {
		return
	}
	ctx.err = err
	close(ctx.done)
}

type deadlineCtx struct {
	*cancelCtx
	deadline time.Time
}

func (ctx *deadlineCtx) Deadline() (deadline time.Time, ok bool) {
	return ctx.deadline, true
}

var DeadlineExceeded = deadlineExceededErr{}

type deadlineExceededErr struct{}

func (deadlineExceededErr) Error() string   { return "deadline exceeded " }
func (deadlineExceededErr) Timeout() bool   { return true }
func (deadlineExceededErr) Temporary() bool { return true }

func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc) {
	cctx, cancel := WithCancel(parent)

	ctx := &deadlineCtx{
		cancelCtx: cctx.(*cancelCtx),
		deadline:  deadline,
	}

	t := time.AfterFunc(time.Until(deadline), func() {
		ctx.cancel(DeadlineExceeded)
	})

	stop := func() {
		t.Stop()
		cancel()
	}

	return ctx, stop
}

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
	return WithDeadline(parent, time.Now().Add(timeout))
}

type valueCtx struct {
	Context
	value, key interface{}
}

func (ctx *valueCtx) Value(key interface{}) interface{} {
	if key == ctx.key {
		return ctx.value
	}
	return ctx.Context.Value(key)
}

func WithValue(parent Context, key, value interface{}) Context {
	if key == nil {
		panic("key is nil")
	}
	if !reflect.TypeOf(key).Comparable() {
		panic("key is not comparable")
	}
	return &valueCtx{
		Context: parent,
		key:     key,
		value:   value,
	}
}
<commit_msg>replaced time.Until with go1.7 compatible equivalent<commit_after>// Copyright 2016 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 writing, software distributed
// under the License is distributed on a "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 contextimpl

import (
	"errors"
	"reflect"
	"sync"
	"time"
)

type Context interface {
	Deadline() (deadline time.Time, ok bool)
	Done() <-chan struct{}
	Err() error
	Value(key interface{}) interface{}
}

type emptyCtx int

func (emptyCtx) Deadline() (deadline time.Time, ok bool) { return }
func (emptyCtx) Done() <-chan struct{}                   { return nil }
func (emptyCtx) Err() error                              { return nil }
func (emptyCtx) Value(key interface{}) interface{}       { return nil }

var (
	background = new(emptyCtx)
	todo       = new(emptyCtx)
)

func Background() Context { return background }
func TODO() Context       { return todo }

type cancelCtx struct {
	Context
	done chan struct{}
	err  error
	mu   sync.Mutex
}

func (ctx *cancelCtx) Done() <-chan struct{} { return ctx.done }
func (ctx *cancelCtx) Err() error {
	ctx.mu.Lock()
	defer ctx.mu.Unlock()
	return ctx.err
}

var Canceled = errors.New("context canceled")

type CancelFunc func()

func WithCancel(parent Context) (Context, CancelFunc) {
	ctx := &cancelCtx{
		Context: parent,
		done:    make(chan struct{}),
	}

	cancel := func() { ctx.cancel(Canceled) }

	go func() {
		select {
		case <-parent.Done():
			ctx.cancel(parent.Err())
		case <-ctx.Done():
		}
	}()

	return ctx, cancel
}

func (ctx *cancelCtx) cancel(err error) {
	ctx.mu.Lock()
	defer ctx.mu.Unlock()
	if ctx.err != nil {
		return
	}
	ctx.err = err
	close(ctx.done)
}

type deadlineCtx struct {
	*cancelCtx
	deadline time.Time
}

func (ctx *deadlineCtx) Deadline() (deadline time.Time, ok bool) {
	return ctx.deadline, true
}

var DeadlineExceeded = deadlineExceededErr{}

type deadlineExceededErr struct{}

func (deadlineExceededErr) Error() string   { return "deadline exceeded " }
func (deadlineExceededErr) Timeout() bool   { return true }
func (deadlineExceededErr) Temporary() bool { return true }

func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc) {
	cctx, cancel := WithCancel(parent)

	ctx := &deadlineCtx{
		cancelCtx: cctx.(*cancelCtx),
		deadline:  deadline,
	}

	// the line below has been replaced to work with version previous to go 1.8.
	// timeout := time.Until(deadline)
	timeout := deadline.Sub(time.Now())
	t := time.AfterFunc(timeout, func() {
		ctx.cancel(DeadlineExceeded)
	})

	stop := func() {
		t.Stop()
		cancel()
	}

	return ctx, stop
}

func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
	return WithDeadline(parent, time.Now().Add(timeout))
}

type valueCtx struct {
	Context
	value, key interface{}
}

func (ctx *valueCtx) Value(key interface{}) interface{} {
	if key == ctx.key {
		return ctx.value
	}
	return ctx.Context.Value(key)
}

func WithValue(parent Context, key, value interface{}) Context {
	if key == nil {
		panic("key is nil")
	}
	if !reflect.TypeOf(key).Comparable() {
		panic("key is not comparable")
	}
	return &valueCtx{
		Context: parent,
		key:     key,
		value:   value,
	}
}
<|endoftext|>