<commit_before>package gui

import (
  "glop/gin"
  "gl"
)

type Point struct {
  X, Y int
}

func (p Point) Add(q Point) Point {
  return Point{
    X: p.X + q.X,
    Y: p.Y + q.Y,
  }
}
func (p Point) Inside(r Region) bool {
  if p.X < r.X {
    return false
  }
  if p.Y < r.Y {
    return false
  }
  if p.X > r.X+r.Dx {
    return false
  }
  if p.Y > r.Y+r.Dy {
    return false
  }
  return true
}

type Dims struct {
  Dx, Dy int
}
type Region struct {
  Point
  Dims
}

func (r Region) Add(p Point) Region {
  return Region{
    r.Point.Add(p),
    r.Dims,
  }
}

// Returns a region that is no larger than r that fits inside t.  The region
// will be located as closely as possible to r.
func (r Region) Fit(t Region) Region {
  if r.Dx > t.Dx {
    r.Dx = t.Dx
  }
  if r.Dy > t.Dy {
    r.Dy = t.Dy
  }
  if r.X < t.X {
    r.X = t.X
  }
  if r.X + r.Dx > t.X + t.Dx {
    r.X -= (r.X + r.Dx) - (t.X + t.Dx)
  }
  if r.Y + r.Dy > t.Y + t.Dy {
    r.Y -= (r.Y + r.Dy) - (t.Y + t.Dy)
  }
  return r
}

func (r Region) Isect(s Region) Region {
  if r.X < s.X {
    r.Dx -= s.X - r.X
    r.X = s.X
  }
  if r.Y < s.Y {
    r.Dy -= s.Y - r.Y
    r.Y = s.Y
  }
  if r.X + r.Dx > s.X + s.Dx {
    r.Dx -= (r.X + r.Dx) - (s.X + s.Dx)
  }
  if r.Y + r.Dy > s.Y + s.Dy {
    r.Dy -= (r.Y + r.Dy) - (s.Y + s.Dy)
  }
  return r
}

func (r Region) Size() int {
  return r.Dx * r.Dy
}

// Need a global stack of regions because opengl only handles pushing/popping
// the state of the enable bits for each clip plane, not the planes themselves
var clippers []Region

func (r Region) setClipPlanes() {
  var eqs [][4]float64
  eqs = append(eqs, [4]float64{1, 0, 0, -float64(r.X)})
  eqs = append(eqs, [4]float64{-1, 0, 0, float64(r.X + r.Dx)})
  eqs = append(eqs, [4]float64{0, 1, 0, -float64(r.Y)})
  eqs = append(eqs, [4]float64{0, -1, 0, float64(r.Y + r.Dy)})
  gl.ClipPlane(gl.CLIP_PLANE0, &eqs[0][0])
  gl.ClipPlane(gl.CLIP_PLANE1, &eqs[1][0])
  gl.ClipPlane(gl.CLIP_PLANE2, &eqs[2][0])
  gl.ClipPlane(gl.CLIP_PLANE3, &eqs[3][0])
}

func (r Region) PushClipPlanes() {
  if len(clippers) == 0 {
    gl.Enable(gl.CLIP_PLANE0)
    gl.Enable(gl.CLIP_PLANE1)
    gl.Enable(gl.CLIP_PLANE2)
    gl.Enable(gl.CLIP_PLANE3)
    r.setClipPlanes()
    clippers = append(clippers, r)
  } else {
    cur := clippers[len(clippers)-1]
    clippers = append(clippers, r.Isect(cur))
  }
}
func (r Region) PopClipPlanes() {
  clippers = clippers[0 : len(clippers)-1]
  if len(clippers) == 0 {
    gl.Disable(gl.CLIP_PLANE0)
    gl.Disable(gl.CLIP_PLANE1)
    gl.Disable(gl.CLIP_PLANE2)
    gl.Disable(gl.CLIP_PLANE3)
  } else {
    clippers[len(clippers)-1].setClipPlanes()
  }
}

//func (r Region) setViewport() {
//  gl.Viewport(r.Point.X, r.Point.Y, r.Dims.Dx, r.Dims.Dy)
//}

type Zone interface {
  // Returns the dimensions that this Widget would like available to
  // render itself.  A Widget should only update the value it returns from
  // this method when its Think() method is called.
  Requested() Dims

  // Returns ex,ey, where ex and ey indicate whether this Widget is
  // capable of expanding along the X and Y axes, respectively.
  Expandable() (bool, bool)

  // Returns the region that this Widget used to render itself the last
  // time it was rendered.  Should be completely contained within the
  // region that was passed to it on its last call to Render.
  Rendered() Region
}

type EventGroup struct {
  gin.EventGroup
  Focus bool
}

type Widget interface {
  Zone

  // Called regularly with a timestamp and a function that checkes whether a Widget is
  // the widget that currently has focus
  Think(*Gui, int64)

  // Returns true if this widget or any of its children consumed the
  // event group
  Respond(*Gui, EventGroup) bool

  Draw(Region)
  DrawFocused(Region)
  String() string
}
type CoreWidget interface {
  DoThink(int64, bool)

  // If change_focus is true, then the EventGroup will be consumed,
  // regardless of the value of consume
  DoRespond(EventGroup) (consume, change_focus bool)
  Zone

  Draw(Region)
  DrawFocused(Region)

  GetChildren() []Widget
  String() string
}
type EmbeddedWidget interface {
  Think(*Gui, int64)
  Respond(*Gui, EventGroup) (consume bool)
}
type BasicWidget struct {
  CoreWidget
}

func (w *BasicWidget) Think(gui *Gui, t int64) {
  kids := w.GetChildren()
  for i := range kids {
    kids[i].Think(gui, t)
  }
  w.DoThink(t, w == gui.FocusWidget())
}
func (w *BasicWidget) Respond(gui *Gui, event_group EventGroup) bool {
  cursor := event_group.Events[0].Key.Cursor()
  if cursor != nil {
    var p Point
    p.X, p.Y = cursor.Point()
    if !p.Inside(w.Rendered()) {
      return false
    }
  }
  consume, change_focus := w.DoRespond(event_group)

  if change_focus {
    if event_group.Focus {
      gui.DropFocus()
    } else {
      gui.TakeFocus(w)
    }
    return true
  }
  if consume {
    return true
  }

  kids := w.GetChildren()
  for i := len(kids) - 1; i >= 0; i-- {
    if kids[i].Respond(gui, event_group) {
      return true
    }
  }
  return false
}

type BasicZone struct {
  Request_dims  Dims
  Render_region Region
  Ex, Ey        bool
}

func (bz BasicZone) Requested() Dims {
  return bz.Request_dims
}
func (bz BasicZone) Rendered() Region {
  return bz.Render_region
}
func (bz BasicZone) Expandable() (bool, bool) {
  return bz.Ex, bz.Ey
}

type CollapsableZone struct {
  Collapsed     bool
  Request_dims  Dims
  Render_region Region
  Ex, Ey        bool
}

func (cz CollapsableZone) Requested() Dims {
  if cz.Collapsed {
    return Dims{}
  }
  return cz.Request_dims
}
func (cz CollapsableZone) Rendered() Region {
  if cz.Collapsed {
    return Region{Point: cz.Render_region.Point}
  }
  return cz.Render_region
}
func (cz *CollapsableZone) Expandable() (bool, bool) {
  if cz.Collapsed {
    return false, false
  }
  return cz.Ex, cz.Ey
}

// Embed a Clickable object to run a specified function when the widget
// is clicked and run a specified function.
type Clickable struct {
  on_click func(int64)
}

func (c Clickable) DoRespond(event_group EventGroup) (bool, bool) {
  event := event_group.Events[0]
  if event.Type == gin.Press && event.Key.Id() == gin.MouseLButton {
    c.on_click(event_group.Timestamp)
    return true, false
  }
  return false, false
}

type NonFocuser struct{}

func (n NonFocuser) DrawFocused(Region) {}

type NonThinker struct{}

func (n NonThinker) DoThink(int64, bool) {}

type NonResponder struct{}

func (n NonResponder) DoRespond(EventGroup) (bool, bool) {
  return false, false
}

type Childless struct{}

func (c Childless) GetChildren() []Widget { return nil }

// Wrappers are used to wrap existing widgets inside another widget to add some
// specific behavior (like making it hideable).  This can also be done by creating
// a new widget and embedding the appropriate structs, but sometimes this is more
// convenient.
type Wrapper struct {
  Child Widget
}

func (w Wrapper) GetChildren() []Widget { return []Widget{w.Child} }
func (w Wrapper) Draw(region Region) {
  w.Child.Draw(region)
}

type StandardParent struct {
  Children []Widget
}

func (s *StandardParent) GetChildren() []Widget {
  return s.Children
}
func (s *StandardParent) AddChild(w Widget) {
  s.Children = append(s.Children, w)
}
func (s *StandardParent) RemoveChild(w Widget) {
  for i := range s.Children {
    if s.Children[i] == w {
      s.Children[i] = s.Children[len(s.Children)-1]
      s.Children = s.Children[0 : len(s.Children)-1]
      return
    }
  }
}
func (s *StandardParent) ReplaceChild(old, new Widget) {
  for i := range s.Children {
    if s.Children[i] == old {
      s.Children[i] = new
      return
    }
  }
}
func (s *StandardParent) RemoveAllChildren() {
  s.Children = s.Children[0:0]
}

type rootWidget struct {
  EmbeddedWidget
  StandardParent
  BasicZone
  NonResponder
  NonThinker
  NonFocuser
}

func (r *rootWidget) String() string {
  return "root"
}

func (r *rootWidget) Draw(region Region) {
  r.Render_region = region
  for i := range r.Children {
    r.Children[i].Draw(region)
  }
}

type Gui struct {
  root rootWidget

  // Stack of widgets that have focus
  focus []Widget
}

func Make(dispatcher gin.EventDispatcher, dims Dims, font_path string) (*Gui, error) {
  err := LoadFontAs(font_path, "standard")
  if err != nil {
    return nil, err
  }
  var g Gui
  g.root.EmbeddedWidget = &BasicWidget{CoreWidget: &g.root}
  g.root.Request_dims = dims
  g.root.Render_region.Dims = dims
  dispatcher.RegisterEventListener(&g)
  return &g, nil
}

func (g *Gui) Draw() {
  gl.MatrixMode(gl.PROJECTION)
  gl.LoadIdentity()
  region := g.root.Render_region
  gl.Ortho(float64(region.X), float64(region.X+region.Dx), float64(region.Y), float64(region.Y+region.Dy), 1000, -1000)
  gl.ClearColor(0, 0, 0, 1)
  gl.Clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT)
  gl.MatrixMode(gl.MODELVIEW)
  gl.LoadIdentity()
  g.root.Draw(region)
  if g.FocusWidget() != nil {
    g.FocusWidget().DrawFocused(region)
  }
}

// TODO: Shouldn't be exposing this
func (g *Gui) Think(t int64) {
  g.root.Think(g, t)
}

// TODO: Shouldn't be exposing this
func (g *Gui) HandleEventGroup(gin_group gin.EventGroup) {
  event_group := EventGroup{gin_group, false}
  if len(g.focus) > 0 {
    event_group.Focus = true
    consume := g.focus[len(g.focus)-1].Respond(g, event_group)
    if consume {
      return
    }
    event_group.Focus = false
  }
  g.root.Respond(g, event_group)
}

func (g *Gui) AddChild(w Widget) {
  g.root.AddChild(w)
}

func (g *Gui) RemoveChild(w Widget) {
  g.root.RemoveChild(w)
}

func (g *Gui) TakeFocus(w Widget) {
  if len(g.focus) == 0 {
    g.focus = append(g.focus, nil)
  }
  g.focus[len(g.focus)-1] = w
}

func (g *Gui) DropFocus() {
  g.focus = g.focus[0 : len(g.focus)-1]
}

func (g *Gui) FocusWidget() Widget {
  if len(g.focus) == 0 {
    return nil
  }
  return g.focus[len(g.focus)-1]
}
<commit_msg>Fixed a bug with clip planes<commit_after>package gui

import (
  "glop/gin"
  "gl"
)

type Point struct {
  X, Y int
}

func (p Point) Add(q Point) Point {
  return Point{
    X: p.X + q.X,
    Y: p.Y + q.Y,
  }
}
func (p Point) Inside(r Region) bool {
  if p.X < r.X {
    return false
  }
  if p.Y < r.Y {
    return false
  }
  if p.X > r.X+r.Dx {
    return false
  }
  if p.Y > r.Y+r.Dy {
    return false
  }
  return true
}

type Dims struct {
  Dx, Dy int
}
type Region struct {
  Point
  Dims
}

func (r Region) Add(p Point) Region {
  return Region{
    r.Point.Add(p),
    r.Dims,
  }
}

// Returns a region that is no larger than r that fits inside t.  The region
// will be located as closely as possible to r.
func (r Region) Fit(t Region) Region {
  if r.Dx > t.Dx {
    r.Dx = t.Dx
  }
  if r.Dy > t.Dy {
    r.Dy = t.Dy
  }
  if r.X < t.X {
    r.X = t.X
  }
  if r.X + r.Dx > t.X + t.Dx {
    r.X -= (r.X + r.Dx) - (t.X + t.Dx)
  }
  if r.Y + r.Dy > t.Y + t.Dy {
    r.Y -= (r.Y + r.Dy) - (t.Y + t.Dy)
  }
  return r
}

func (r Region) Isect(s Region) Region {
  if r.X < s.X {
    r.Dx -= s.X - r.X
    r.X = s.X
  }
  if r.Y < s.Y {
    r.Dy -= s.Y - r.Y
    r.Y = s.Y
  }
  if r.X + r.Dx > s.X + s.Dx {
    r.Dx -= (r.X + r.Dx) - (s.X + s.Dx)
  }
  if r.Y + r.Dy > s.Y + s.Dy {
    r.Dy -= (r.Y + r.Dy) - (s.Y + s.Dy)
  }
  if r.Dx < 0 { r.Dx = 0 }
  if r.Dy < 0 { r.Dy = 0 }
  return r
}

func (r Region) Size() int {
  return r.Dx * r.Dy
}

// Need a global stack of regions because opengl only handles pushing/popping
// the state of the enable bits for each clip plane, not the planes themselves
var clippers []Region

func (r Region) setClipPlanes() {
  var eqs [][4]float64
  eqs = append(eqs, [4]float64{1, 0, 0, -float64(r.X)})
  eqs = append(eqs, [4]float64{-1, 0, 0, float64(r.X + r.Dx)})
  eqs = append(eqs, [4]float64{0, 1, 0, -float64(r.Y)})
  eqs = append(eqs, [4]float64{0, -1, 0, float64(r.Y + r.Dy)})
  gl.ClipPlane(gl.CLIP_PLANE0, &eqs[0][0])
  gl.ClipPlane(gl.CLIP_PLANE1, &eqs[1][0])
  gl.ClipPlane(gl.CLIP_PLANE2, &eqs[2][0])
  gl.ClipPlane(gl.CLIP_PLANE3, &eqs[3][0])
}

func (r Region) PushClipPlanes() {
  if len(clippers) == 0 {
    gl.Enable(gl.CLIP_PLANE0)
    gl.Enable(gl.CLIP_PLANE1)
    gl.Enable(gl.CLIP_PLANE2)
    gl.Enable(gl.CLIP_PLANE3)
    r.setClipPlanes()
    clippers = append(clippers, r)
  } else {
    cur := clippers[len(clippers)-1]
    clippers = append(clippers, r.Isect(cur))
    clippers[len(clippers)-1].setClipPlanes()
  }
}
func (r Region) PopClipPlanes() {
  clippers = clippers[0 : len(clippers)-1]
  if len(clippers) == 0 {
    gl.Disable(gl.CLIP_PLANE0)
    gl.Disable(gl.CLIP_PLANE1)
    gl.Disable(gl.CLIP_PLANE2)
    gl.Disable(gl.CLIP_PLANE3)
  } else {
    clippers[len(clippers)-1].setClipPlanes()
  }
}

//func (r Region) setViewport() {
//  gl.Viewport(r.Point.X, r.Point.Y, r.Dims.Dx, r.Dims.Dy)
//}

type Zone interface {
  // Returns the dimensions that this Widget would like available to
  // render itself.  A Widget should only update the value it returns from
  // this method when its Think() method is called.
  Requested() Dims

  // Returns ex,ey, where ex and ey indicate whether this Widget is
  // capable of expanding along the X and Y axes, respectively.
  Expandable() (bool, bool)

  // Returns the region that this Widget used to render itself the last
  // time it was rendered.  Should be completely contained within the
  // region that was passed to it on its last call to Render.
  Rendered() Region
}

type EventGroup struct {
  gin.EventGroup
  Focus bool
}

type Widget interface {
  Zone

  // Called regularly with a timestamp and a function that checkes whether a Widget is
  // the widget that currently has focus
  Think(*Gui, int64)

  // Returns true if this widget or any of its children consumed the
  // event group
  Respond(*Gui, EventGroup) bool

  Draw(Region)
  DrawFocused(Region)
  String() string
}
type CoreWidget interface {
  DoThink(int64, bool)

  // If change_focus is true, then the EventGroup will be consumed,
  // regardless of the value of consume
  DoRespond(EventGroup) (consume, change_focus bool)
  Zone

  Draw(Region)
  DrawFocused(Region)

  GetChildren() []Widget
  String() string
}
type EmbeddedWidget interface {
  Think(*Gui, int64)
  Respond(*Gui, EventGroup) (consume bool)
}
type BasicWidget struct {
  CoreWidget
}

func (w *BasicWidget) Think(gui *Gui, t int64) {
  kids := w.GetChildren()
  for i := range kids {
    kids[i].Think(gui, t)
  }
  w.DoThink(t, w == gui.FocusWidget())
}
func (w *BasicWidget) Respond(gui *Gui, event_group EventGroup) bool {
  cursor := event_group.Events[0].Key.Cursor()
  if cursor != nil {
    var p Point
    p.X, p.Y = cursor.Point()
    if !p.Inside(w.Rendered()) {
      return false
    }
  }
  consume, change_focus := w.DoRespond(event_group)

  if change_focus {
    if event_group.Focus {
      gui.DropFocus()
    } else {
      gui.TakeFocus(w)
    }
    return true
  }
  if consume {
    return true
  }

  kids := w.GetChildren()
  for i := len(kids) - 1; i >= 0; i-- {
    if kids[i].Respond(gui, event_group) {
      return true
    }
  }
  return false
}

type BasicZone struct {
  Request_dims  Dims
  Render_region Region
  Ex, Ey        bool
}

func (bz BasicZone) Requested() Dims {
  return bz.Request_dims
}
func (bz BasicZone) Rendered() Region {
  return bz.Render_region
}
func (bz BasicZone) Expandable() (bool, bool) {
  return bz.Ex, bz.Ey
}

type CollapsableZone struct {
  Collapsed     bool
  Request_dims  Dims
  Render_region Region
  Ex, Ey        bool
}

func (cz CollapsableZone) Requested() Dims {
  if cz.Collapsed {
    return Dims{}
  }
  return cz.Request_dims
}
func (cz CollapsableZone) Rendered() Region {
  if cz.Collapsed {
    return Region{Point: cz.Render_region.Point}
  }
  return cz.Render_region
}
func (cz *CollapsableZone) Expandable() (bool, bool) {
  if cz.Collapsed {
    return false, false
  }
  return cz.Ex, cz.Ey
}

// Embed a Clickable object to run a specified function when the widget
// is clicked and run a specified function.
type Clickable struct {
  on_click func(int64)
}

func (c Clickable) DoRespond(event_group EventGroup) (bool, bool) {
  event := event_group.Events[0]
  if event.Type == gin.Press && event.Key.Id() == gin.MouseLButton {
    c.on_click(event_group.Timestamp)
    return true, false
  }
  return false, false
}

type NonFocuser struct{}

func (n NonFocuser) DrawFocused(Region) {}

type NonThinker struct{}

func (n NonThinker) DoThink(int64, bool) {}

type NonResponder struct{}

func (n NonResponder) DoRespond(EventGroup) (bool, bool) {
  return false, false
}

type Childless struct{}

func (c Childless) GetChildren() []Widget { return nil }

// Wrappers are used to wrap existing widgets inside another widget to add some
// specific behavior (like making it hideable).  This can also be done by creating
// a new widget and embedding the appropriate structs, but sometimes this is more
// convenient.
type Wrapper struct {
  Child Widget
}

func (w Wrapper) GetChildren() []Widget { return []Widget{w.Child} }
func (w Wrapper) Draw(region Region) {
  w.Child.Draw(region)
}

type StandardParent struct {
  Children []Widget
}

func (s *StandardParent) GetChildren() []Widget {
  return s.Children
}
func (s *StandardParent) AddChild(w Widget) {
  s.Children = append(s.Children, w)
}
func (s *StandardParent) RemoveChild(w Widget) {
  for i := range s.Children {
    if s.Children[i] == w {
      s.Children[i] = s.Children[len(s.Children)-1]
      s.Children = s.Children[0 : len(s.Children)-1]
      return
    }
  }
}
func (s *StandardParent) ReplaceChild(old, new Widget) {
  for i := range s.Children {
    if s.Children[i] == old {
      s.Children[i] = new
      return
    }
  }
}
func (s *StandardParent) RemoveAllChildren() {
  s.Children = s.Children[0:0]
}

type rootWidget struct {
  EmbeddedWidget
  StandardParent
  BasicZone
  NonResponder
  NonThinker
  NonFocuser
}

func (r *rootWidget) String() string {
  return "root"
}

func (r *rootWidget) Draw(region Region) {
  r.Render_region = region
  for i := range r.Children {
    r.Children[i].Draw(region)
  }
}

type Gui struct {
  root rootWidget

  // Stack of widgets that have focus
  focus []Widget
}

func Make(dispatcher gin.EventDispatcher, dims Dims, font_path string) (*Gui, error) {
  err := LoadFontAs(font_path, "standard")
  if err != nil {
    return nil, err
  }
  var g Gui
  g.root.EmbeddedWidget = &BasicWidget{CoreWidget: &g.root}
  g.root.Request_dims = dims
  g.root.Render_region.Dims = dims
  dispatcher.RegisterEventListener(&g)
  return &g, nil
}

func (g *Gui) Draw() {
  gl.MatrixMode(gl.PROJECTION)
  gl.LoadIdentity()
  region := g.root.Render_region
  gl.Ortho(float64(region.X), float64(region.X+region.Dx), float64(region.Y), float64(region.Y+region.Dy), 1000, -1000)
  gl.ClearColor(0, 0, 0, 1)
  gl.Clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT)
  gl.MatrixMode(gl.MODELVIEW)
  gl.LoadIdentity()
  g.root.Draw(region)
  if g.FocusWidget() != nil {
    g.FocusWidget().DrawFocused(region)
  }
}

// TODO: Shouldn't be exposing this
func (g *Gui) Think(t int64) {
  g.root.Think(g, t)
}

// TODO: Shouldn't be exposing this
func (g *Gui) HandleEventGroup(gin_group gin.EventGroup) {
  event_group := EventGroup{gin_group, false}
  if len(g.focus) > 0 {
    event_group.Focus = true
    consume := g.focus[len(g.focus)-1].Respond(g, event_group)
    if consume {
      return
    }
    event_group.Focus = false
  }
  g.root.Respond(g, event_group)
}

func (g *Gui) AddChild(w Widget) {
  g.root.AddChild(w)
}

func (g *Gui) RemoveChild(w Widget) {
  g.root.RemoveChild(w)
}

func (g *Gui) TakeFocus(w Widget) {
  if len(g.focus) == 0 {
    g.focus = append(g.focus, nil)
  }
  g.focus[len(g.focus)-1] = w
}

func (g *Gui) DropFocus() {
  g.focus = g.focus[0 : len(g.focus)-1]
}

func (g *Gui) FocusWidget() Widget {
  if len(g.focus) == 0 {
    return nil
  }
  return g.focus[len(g.focus)-1]
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"sync"

	log "github.com/GameGophers/libs/nsq-logger"
	"github.com/xtaci/go-pubsub"
	"golang.org/x/net/context"
)

import (
	. "proto"
)

const (
	SERVICE = "[CHAT]"
)

const (
	BOLTDB_FILE    = "/data/CHAT.DAT"
	BOLTDB_BUCKET  = "EPS"
	MAX_QUEUE_SIZE = 128 // num of message kept
)

var (
	OK                   = &Chat_Nil{}
	ERROR_ALREADY_EXISTS = errors.New("id already exists")
	ERROR_NOT_EXISTS     = errors.New("id not exists")
)

type EndPoint struct {
	inbox []Chat_Message
	ps    *pubsub.PubSub
	sync.Mutex
}

func (ep *EndPoint) Push(msg *Chat_Message) {
	ep.Lock()
	defer ep.Unlock()
	if len(ep.inbox) > MAX_QUEUE_SIZE {
		ep.inbox = append(ep.inbox[1:], *msg)
	} else {
		ep.inbox = append(ep.inbox, *msg)
	}
}

func (ep *EndPoint) Read() []Chat_Message {
	ep.Lock()
	defer ep.Unlock()
	return append([]Chat_Message(nil), ep.inbox...)
}

func NewEndPoint() *EndPoint {
	u := &EndPoint{}
	u.ps = pubsub.New()
	return u
}

type server struct {
	eps map[uint64]*EndPoint
	sync.RWMutex
}

func (s *server) read_ep(id uint64) *EndPoint {
	s.RLock()
	defer s.RUnlock()
	return s.eps[id]
}

func (s *server) init() {
	s.eps = make(map[uint64]*EndPoint)
}

func (s *server) Subscribe(p *Chat_Id, stream ChatService_SubscribeServer) error {
	die := make(chan bool)
	f := pubsub.NewWrap(func(msg *Chat_Message) {
		if err := stream.Send(msg); err != nil {
			close(die)
		}
	})

	log.Tracef("new subscriber: %p", f)

	ep := s.read_ep(p.Id)
	if ep == nil {
		log.Errorf("cannot find endpoint %v", p)
		return ERROR_NOT_EXISTS
	}

	ep.ps.Sub(f)
	defer func() {
		ep.ps.Leave(f)
	}()

	<-die
	return nil
}

func (s *server) Read(p *Chat_Id, stream ChatService_ReadServer) error {
	ep := s.read_ep(p.Id)
	if ep == nil {
		log.Errorf("cannot find endpoint %v", p)
		return ERROR_NOT_EXISTS
	}

	msgs := ep.Read()
	for k := range msgs {
		if err := stream.Send(&msgs[k]); err != nil {
			return err
		}
	}
	return nil
}

func (s *server) Send(ctx context.Context, msg *Chat_Message) (*Chat_Nil, error) {
	ep := s.read_ep(msg.Id)
	if ep == nil {
		return nil, ERROR_NOT_EXISTS
	}

	ep.ps.Pub(msg)
	ep.Push(msg)
	return OK, nil
}

func (s *server) Reg(ctx context.Context, p *Chat_Id) (*Chat_Nil, error) {
	s.Lock()
	defer s.Unlock()
	ep := s.eps[p.Id]
	if ep != nil {
		log.Errorf("id already exists:%v", p.Id)
		return nil, ERROR_ALREADY_EXISTS
	}

	s.eps[p.Id] = NewEndPoint()
	return OK, nil
}
<commit_msg>add persistence<commit_after>package main

import (
	"errors"
	"fmt"
	"os"
	"os/signal"
	"strconv"
	"sync"
	"syscall"
	"time"

	log "github.com/GameGophers/libs/nsq-logger"
	"github.com/boltdb/bolt"
	"github.com/xtaci/go-pubsub"
	"golang.org/x/net/context"
	"gopkg.in/vmihailenco/msgpack.v2"
)

import (
	. "proto"
)

const (
	SERVICE = "[CHAT]"
)

const (
	BOLTDB_FILE    = "/data/CHAT.DAT"
	BOLTDB_BUCKET  = "EPS"
	MAX_QUEUE_SIZE = 128 // num of message kept
	PENDING_SIZE   = 65536
	CHECK_INTERVAL = time.Minute // if ranking has changed, how long to check
)

var (
	OK                   = &Chat_Nil{}
	ERROR_ALREADY_EXISTS = errors.New("id already exists")
	ERROR_NOT_EXISTS     = errors.New("id not exists")
)

type EndPoint struct {
	inbox []Chat_Message
	ps    *pubsub.PubSub
	sync.Mutex
}

func (ep *EndPoint) Push(msg *Chat_Message) {
	ep.Lock()
	defer ep.Unlock()
	if len(ep.inbox) > MAX_QUEUE_SIZE {
		ep.inbox = append(ep.inbox[1:], *msg)
	} else {
		ep.inbox = append(ep.inbox, *msg)
	}
}

func (ep *EndPoint) Read() []Chat_Message {
	ep.Lock()
	defer ep.Unlock()
	return append([]Chat_Message(nil), ep.inbox...)
}

func NewEndPoint() *EndPoint {
	u := &EndPoint{}
	u.ps = pubsub.New()
	return u
}

type server struct {
	eps     map[uint64]*EndPoint
	pending chan uint64
	sync.RWMutex
}

func (s *server) init() {
	s.eps = make(map[uint64]*EndPoint)
	s.pending = make(chan uint64, PENDING_SIZE)
	s.restore()
	go s.persistence_task()
}

func (s *server) read_ep(id uint64) *EndPoint {
	s.RLock()
	defer s.RUnlock()
	return s.eps[id]
}

func (s *server) Subscribe(p *Chat_Id, stream ChatService_SubscribeServer) error {
	die := make(chan bool)
	f := pubsub.NewWrap(func(msg *Chat_Message) {
		if err := stream.Send(msg); err != nil {
			close(die)
		}
	})

	log.Tracef("new subscriber: %p", f)

	ep := s.read_ep(p.Id)
	if ep == nil {
		log.Errorf("cannot find endpoint %v", p)
		return ERROR_NOT_EXISTS
	}

	ep.ps.Sub(f)
	defer func() {
		ep.ps.Leave(f)
	}()

	<-die
	return nil
}

func (s *server) Read(p *Chat_Id, stream ChatService_ReadServer) error {
	ep := s.read_ep(p.Id)
	if ep == nil {
		log.Errorf("cannot find endpoint %v", p)
		return ERROR_NOT_EXISTS
	}

	msgs := ep.Read()
	for k := range msgs {
		if err := stream.Send(&msgs[k]); err != nil {
			return err
		}
	}
	return nil
}

func (s *server) Send(ctx context.Context, msg *Chat_Message) (*Chat_Nil, error) {
	ep := s.read_ep(msg.Id)
	if ep == nil {
		return nil, ERROR_NOT_EXISTS
	}

	ep.ps.Pub(msg)
	ep.Push(msg)
	s.pending <- msg.Id
	return OK, nil
}

func (s *server) Reg(ctx context.Context, p *Chat_Id) (*Chat_Nil, error) {
	s.Lock()
	defer s.Unlock()
	ep := s.eps[p.Id]
	if ep != nil {
		log.Errorf("id already exists:%v", p.Id)
		return nil, ERROR_ALREADY_EXISTS
	}

	s.eps[p.Id] = NewEndPoint()
	s.pending <- p.Id
	return OK, nil
}

// persistence ranking tree into db
func (s *server) persistence_task() {
	timer := time.After(CHECK_INTERVAL)
	db := s.open_db()
	changes := make(map[uint64]bool)
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGTERM)

	for {
		select {
		case key := <-s.pending:
			changes[key] = true
		case <-timer:
			s.dump(db, changes)
			log.Infof("perisisted %v endpoints:", len(changes))
			changes = make(map[uint64]bool)
			timer = time.After(CHECK_INTERVAL)
		case <-sig:
			s.dump(db, changes)
			db.Close()
			log.Info("SIGTERM")
			os.Exit(0)
		}
	}
}

func (s *server) open_db() *bolt.DB {
	db, err := bolt.Open(BOLTDB_FILE, 0600, nil)
	if err != nil {
		log.Critical(err)
		os.Exit(-1)
	}
	// create bulket
	db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte(BOLTDB_BUCKET))
		if err != nil {
			log.Criticalf("create bucket: %s", err)
			os.Exit(-1)
		}
		return nil
	})
	return db
}

func (s *server) dump(db *bolt.DB, changes map[uint64]bool) {
	for k := range changes {
		ep := s.read_ep(k)
		if ep == nil {
			log.Errorf("cannot find endpoint %v", k)
			continue
		}

		// serialization and save
		bin, err := msgpack.Marshal(ep.Read())
		if err != nil {
			log.Critical("cannot marshal:", err)
			continue
		}

		db.Update(func(tx *bolt.Tx) error {
			b := tx.Bucket([]byte(BOLTDB_BUCKET))
			err := b.Put([]byte(fmt.Sprint(k)), bin)
			return err
		})
	}
}

func (s *server) restore() {
	// restore data from db file
	db := s.open_db()
	defer db.Close()
	count := 0
	db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(BOLTDB_BUCKET))
		c := b.Cursor()

		for k, v := c.First(); k != nil; k, v = c.Next() {
			var msg []Chat_Message
			err := msgpack.Unmarshal(v, &msg)
			if err != nil {
				log.Critical("chat data corrupted:", err)
				os.Exit(-1)
			}
			id, err := strconv.ParseUint(string(k), 0, 64)
			if err != nil {
				log.Critical("chat data corrupted:", err)
				os.Exit(-1)
			}
			ep := NewEndPoint()
			ep.inbox = msg
			s.eps[id] = ep
			count++
		}

		return nil
	})

	log.Infof("restored %v chats", count)
}
<|endoftext|>


<commit_before>// Package queue implements a First In First Out (FIFO):
// +-----+           +-----+
// |     |           |     |
// |  3  +---+       |  3  |
// |     |   |       |     |
// +-----+   |       +-----+
//        +--v--+    +-----+    +-----+
//        |     |    |     |    |     |
//        |  2  |    |  2  |    |  3  |
//        |     |    |     |    |     |
//        +-----+    +-----+    +-----+
//        +-----+    +-----+    +-----+
//        |     |    |     |    |     |
//        |  1  |    |  1  |    |  2  |
//        |     |    |     |    |     |
//        +-----+    +--+--+    +-----+
//                      |
//                      v
package queue

// Element is an element in a single linked list
type Element struct {
	prev *Element
	// The value stored in this Element
	Value interface{}
}

// Queue is an abstract datatype, retrieving the items in a chronological order
// from their insertion
type Queue struct {
	top    *Element
	bottom *Element
}

// New create a new Queue structure
func New() *Queue {
	q := new(Queue)
	return q
}

// Enqueue add an item in the top of the Queue
func (q *Queue) Enqueue(value interface{}) {
	e := new(Element)
	e.Value = value
	q.insert(e)
}

// Dequeue retrieve and remove the oldest item in the Queue
func (q *Queue) Dequeue() *Element {
	return q.remove()
}

// IsEmpty define if the Queue is empty
func (q *Queue) IsEmpty() bool {
	return q.bottom == nil
}

// Peek retrieve the oldest item in the Queue
func (q *Queue) Peek() *Element {
	return q.bottom
}

func (q *Queue) insert(e *Element) {
	if q.top != nil {
		q.top.prev = e
	}
	q.top = e
	if q.bottom == nil {
		q.bottom = e
	}
}

func (q *Queue) remove() *Element {
	previously := q.bottom
	if q.bottom != nil {
		q.bottom = q.bottom.prev
	}
	if q.bottom == nil {
		q.top = nil
	}
	return previously
}
<commit_msg>Reduce the size of the diagram once again<commit_after>// Package queue implements a First In First Out (FIFO):
//
// +-----+           +-----+
// |     |           |     |
// |  2  +---+       |  2  |
// |     |   |       |     |
// +-----+   |       +-----+
//        +--v--+    +-----+    +-----+
//        |     |    |     |    |     |
//        |  1  |    |  1  |    |  2  |
//        |     |    |     |    |     |
//        +-----+    +-----+    +-----+
//                      |
//                      v
package queue

// Element is an element in a single linked list
type Element struct {
	prev *Element
	// The value stored in this Element
	Value interface{}
}

// Queue is an abstract datatype, retrieving the items in a chronological order
// from their insertion
type Queue struct {
	top    *Element
	bottom *Element
}

// New create a new Queue structure
func New() *Queue {
	q := new(Queue)
	return q
}

// Enqueue add an item in the top of the Queue
func (q *Queue) Enqueue(value interface{}) {
	e := new(Element)
	e.Value = value
	q.insert(e)
}

// Dequeue retrieve and remove the oldest item in the Queue
func (q *Queue) Dequeue() *Element {
	return q.remove()
}

// IsEmpty define if the Queue is empty
func (q *Queue) IsEmpty() bool {
	return q.bottom == nil
}

// Peek retrieve the oldest item in the Queue
func (q *Queue) Peek() *Element {
	return q.bottom
}

func (q *Queue) insert(e *Element) {
	if q.top != nil {
		q.top.prev = e
	}
	q.top = e
	if q.bottom == nil {
		q.bottom = e
	}
}

func (q *Queue) remove() *Element {
	previously := q.bottom
	if q.bottom != nil {
		q.bottom = q.bottom.prev
	}
	if q.bottom == nil {
		q.top = nil
	}
	return previously
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/md5"
	"encoding/json"
	"errors"
	"fmt"
	"image"
	"image/jpeg"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"os"
	"time"
	"vip/fetch"

	"github.com/golang/groupcache"
	"github.com/gorilla/mux"
)

type UploadResponse struct {
	Url string `json:"url"`
}

type ErrorResponse struct {
	Msg string `json:"error"`
}

type Uploadable struct {
	Data   io.Reader
	Key    string
	Length int64
}

type verifyAuth func(http.ResponseWriter, *http.Request)

func (h verifyAuth) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Enable cross-origin requests
	if domain := os.Getenv("ALLOWED_ORIGIN"); domain != "" {
		if origin := r.Header.Get("Origin"); origin == domain {
			w.Header().Set("Access-Control-Allow-Origin", origin)
			w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
			w.Header().Set("Access-Control-Allow-Headers",
				"Accept, Content-Type, Content-Length, Accept-Encoding, X-Vip-Token, Authorization")
		}
	} else {
		auth := r.Header.Get("X-Vip-Token")
		if auth != authToken {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
	}

	if r.Method == "OPTIONS" {
		return
	}

	h(w, r)
}

func fileKey(bucket string, width int, height int) string {
	seed := rand.New(rand.NewSource(time.Now().UnixNano()))
	key := fmt.Sprintf("%d-%s-%d", seed.Int63(), bucket, time.Now().UnixNano())

	hash := md5.New()
	io.WriteString(hash, key)
	return fmt.Sprintf("%x-%dx%d", hash.Sum(nil), width, height)
}

func handleImageRequest(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		w.WriteHeader(http.StatusMethodNotAllowed)
	}

	w.Header().Set("Cache-Control", "public, max-age=31536000")

	// Client is checking for a cached URI, assume it is valid
	// and return a 304
	if r.Header.Get("If-Modified-Since") != "" {
		w.WriteHeader(http.StatusNotModified)
		return
	}

	gc := fetch.RequestContext(r)

	var data []byte
	err := cache.Get(gc, gc.CacheKey(), groupcache.AllocatingByteSliceSink(&data))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

	w.Header().Set("Content-Type", http.DetectContentType(data))
	http.ServeContent(w, r, gc.ImageId, time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC), bytes.NewReader(data))
}

func handleUpload(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}

	vars := mux.Vars(r)
	bucket := vars["bucket_id"]
	// Set a hard limit in MB on files
	var limit int64 = 5
	if r.ContentLength > limit<<20 {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusRequestEntityTooLarge)
		json.NewEncoder(w).Encode(ErrorResponse{
			Msg: fmt.Sprintf("The file size limit is %dMB", limit),
		})
		return
	} else if r.ContentLength == 0 {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusBadRequest)
		json.NewEncoder(w).Encode(ErrorResponse{
			Msg: fmt.Sprintf("File must have size greater than 0"),
		})
		return
	}

	mime := r.Header.Get("Content-Type")

	data, err := processFile(r.Body, mime, bucket)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
	r.Body.Close()

	err = storage.PutReader(bucket, data.Key, data.Data,
		data.Length, r.Header.Get("Content-Type"))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	uri := r.URL

	if r.URL.Host == "" {
		uri.Host = os.Getenv("URI_HOSTNAME")
		if secure {
			uri.Scheme = "https"
		} else {
			uri.Scheme = "http"
		}
	}

	uri.Path = fmt.Sprintf("%s/%s", bucket, data.Key)

	w.Header().Set("Content-Type", "application/json")

	w.WriteHeader(http.StatusCreated)
	json.NewEncoder(w).Encode(UploadResponse{
		Url: uri.String(),
	})
}

func handlePing(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusOK)
	fmt.Fprintf(w, "pong")
}

func processFile(src io.Reader, mime string, bucket string) (*Uploadable, error) {
	if mime == "image/jpeg" || mime == "image/jpg" {
		image, format, err := fetch.GetRotatedImage(src)
		if err != nil {
			return nil, err
		}
		if format != "jpeg" {
			return nil, errors.New("You sent a bad JPEG file.")
		}

		width := image.Bounds().Size().X
		height := image.Bounds().Size().Y
		key := fileKey(bucket, width, height)

		data := new(bytes.Buffer)
		err = jpeg.Encode(data, image, nil)
		if err != nil {
			return nil, err
		}
		length := int64(data.Len())

		return &Uploadable{data, key, length}, nil

	} else {
		raw, err := ioutil.ReadAll(src)
		if err != nil {
			return nil, err
		}

		data := bytes.NewReader(raw)
		length := int64(data.Len())
		image, _, err := image.Decode(data)
		if err != nil {
			return nil, err
		}

		width := image.Bounds().Size().X
		height := image.Bounds().Size().Y
		key := fileKey(bucket, width, height)

		data.Seek(0, 0)

		upload := Uploadable{data, key, length}
		return &upload, nil
	}
}
<commit_msg>For consistency's sake<commit_after>package main

import (
	"bytes"
	"crypto/md5"
	"encoding/json"
	"errors"
	"fmt"
	"image"
	"image/jpeg"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"os"
	"time"
	"vip/fetch"

	"github.com/golang/groupcache"
	"github.com/gorilla/mux"
)

type UploadResponse struct {
	Url string `json:"url"`
}

type ErrorResponse struct {
	Msg string `json:"error"`
}

type Uploadable struct {
	Data   io.Reader
	Key    string
	Length int64
}

type verifyAuth func(http.ResponseWriter, *http.Request)

func (h verifyAuth) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Enable cross-origin requests
	if domain := os.Getenv("ALLOWED_ORIGIN"); domain != "" {
		if origin := r.Header.Get("Origin"); origin == domain {
			w.Header().Set("Access-Control-Allow-Origin", origin)
			w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
			w.Header().Set("Access-Control-Allow-Headers",
				"Accept, Content-Type, Content-Length, Accept-Encoding, X-Vip-Token, Authorization")
		}
	} else {
		auth := r.Header.Get("X-Vip-Token")
		if auth != authToken {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
	}

	if r.Method == "OPTIONS" {
		return
	}

	h(w, r)
}

func fileKey(bucket string, width int, height int) string {
	seed := rand.New(rand.NewSource(time.Now().UnixNano()))
	key := fmt.Sprintf("%d-%s-%d", seed.Int63(), bucket, time.Now().UnixNano())

	hash := md5.New()
	io.WriteString(hash, key)
	return fmt.Sprintf("%x-%dx%d", hash.Sum(nil), width, height)
}

func handleImageRequest(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		w.WriteHeader(http.StatusMethodNotAllowed)
	}

	w.Header().Set("Cache-Control", "public, max-age=31536000")

	// Client is checking for a cached URI, assume it is valid
	// and return a 304
	if r.Header.Get("If-Modified-Since") != "" {
		w.WriteHeader(http.StatusNotModified)
		return
	}

	gc := fetch.RequestContext(r)

	var data []byte
	err := cache.Get(gc, gc.CacheKey(), groupcache.AllocatingByteSliceSink(&data))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

	w.Header().Set("Content-Type", http.DetectContentType(data))
	http.ServeContent(w, r, gc.ImageId, time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC), bytes.NewReader(data))
}

func handleUpload(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		w.WriteHeader(http.StatusMethodNotAllowed)
		return
	}

	vars := mux.Vars(r)
	bucket := vars["bucket_id"]
	// Set a hard limit in MB on files
	var limit int64 = 5
	if r.ContentLength > limit<<20 {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusRequestEntityTooLarge)
		json.NewEncoder(w).Encode(ErrorResponse{
			Msg: fmt.Sprintf("The file size limit is %dMB", limit),
		})
		return
	} else if r.ContentLength == 0 {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusBadRequest)
		json.NewEncoder(w).Encode(ErrorResponse{
			Msg: fmt.Sprintf("File must have size greater than 0"),
		})
		return
	}

	mime := r.Header.Get("Content-Type")

	data, err := processFile(r.Body, mime, bucket)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
	r.Body.Close()

	err = storage.PutReader(bucket, data.Key, data.Data,
		data.Length, r.Header.Get("Content-Type"))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	uri := r.URL

	if r.URL.Host == "" {
		uri.Host = os.Getenv("URI_HOSTNAME")
		if secure {
			uri.Scheme = "https"
		} else {
			uri.Scheme = "http"
		}
	}

	uri.Path = fmt.Sprintf("%s/%s", bucket, data.Key)

	w.Header().Set("Content-Type", "application/json")

	w.WriteHeader(http.StatusCreated)
	json.NewEncoder(w).Encode(UploadResponse{
		Url: uri.String(),
	})
}

func handlePing(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusOK)
	fmt.Fprintf(w, "pong")
}

func processFile(src io.Reader, mime string, bucket string) (*Uploadable, error) {
	if mime == "image/jpeg" || mime == "image/jpg" {
		image, format, err := fetch.GetRotatedImage(src)
		if err != nil {
			return nil, err
		}
		if format != "jpeg" {
			return nil, errors.New("You sent a bad JPEG file.")
		}

		width := image.Bounds().Size().X
		height := image.Bounds().Size().Y
		key := fileKey(bucket, width, height)

		data := new(bytes.Buffer)
		err = jpeg.Encode(data, image, nil)
		if err != nil {
			return nil, err
		}
		length := int64(data.Len())

		return &Uploadable{data, key, length}, nil

	} else {
		raw, err := ioutil.ReadAll(src)
		if err != nil {
			return nil, err
		}

		data := bytes.NewReader(raw)
		length := int64(data.Len())
		image, _, err := image.Decode(data)
		if err != nil {
			return nil, err
		}

		width := image.Bounds().Size().X
		height := image.Bounds().Size().Y
		key := fileKey(bucket, width, height)

		data.Seek(0, 0)

		return &Uploadable{data, key, length}, nil
	}
}
<|endoftext|>


<commit_before>package main

import (
	"html/template"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/ajg/form"
)

// Fields is a standard handler that pulls the first folder of the response, and
//  lists that topic or author. If there is none, it falls back to listing all
//  topics or authors with the fallback template.
type Fields struct {
	c ServerSection
	i *Index
}

func (h Fields) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	fields := strings.SplitN(r.URL.Path, "/", 3)

	if len(fields) < 3 || fields[1] == "" {
		// to do if a field was not given
		h.i.FallbackSearchResponse(w, h.c.FallbackTemplate)
		return
	} else {
		// to be done if a field was given
		results, err := h.i.ListAllField(h.c.Default, fields[0], 100, 1)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		err = allTemplates.ExecuteTemplate(w, h.c.Template, results)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
		}
	}
}

// FuzzySearch is a normal search format - it should provide a point and click
//  interface to allow searching.
type FuzzySearch struct {
	c ServerSection
	i *Index
}

func (h FuzzySearch) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		// to do if a field was not given
		h.i.FallbackSearchResponse(w, h.c.FallbackTemplate)
		return
	}

	err := r.ParseForm()
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	var values FuzzySearchValues
	err = form.DecodeValues(&values, r.Form)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	results, err := h.i.FuzzySearch(values)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	err = allTemplates.ExecuteTemplate(w, h.c.Template, results)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// QuerySearch is a handler that uses a custom search format to do custom queries.
type QuerySearch struct {
	c ServerSection
	i *Index
}

func (h QuerySearch) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	values := struct {
		s        string `form:"s"`
		page     int    `form:"page"`
		pageSize int    `form:"pageSize"`
	}{}

	if r.Method != http.MethodPost {
		// to do if a field was not given
		h.i.FallbackSearchResponse(w, h.c.FallbackTemplate)
		return
	}

	err := r.ParseForm()
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	err = form.DecodeValues(&values, r.Form)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if values.s == "" {
		// to do if a field was not given
		h.i.FallbackSearchResponse(w, h.c.FallbackTemplate)
		return
	}

	results, err := h.i.QuerySearch(values.s, values.page, values.pageSize)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	err = allTemplates.ExecuteTemplate(w, h.c.Template, results)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// Page is a standard data structure used to render markdown pages.
type Page struct {
	Title    string
	ToC      template.HTML
	Body     template.HTML
	Topics   []string
	Keywords []string
	Authors  []string
}

// Markdown is an http.Handler that renders a markdown file and serves it back.
//  Author and Topic tags before the first major title are parsed and displayed.
//  It is possible to restrict access to a page based on topic tag.
type Markdown struct {
	c ServerSection
}

func (h Markdown) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// If the request is empty, set it to the default.
	if r.URL.Path == "/" {
		r.URL.Path = path.Clean(h.c.Default)
	}

	// If the request doesn't end in .md, add that
	if path.Ext(r.URL.Path) != "md" {
		r.URL.Path = r.URL.Path + ".md"
	}

	pdata := new(PageMetadata)
	err := pdata.LoadPage(h.c.Path + r.URL.Path)
	if err != nil {
		log.Printf("request [ %s ] points bad file target [ %s ] sent to server",
			r.URL.Path, h.c.Path)
		http.Error(w, "Page not Found", http.StatusNotFound)
		return
	}

	if pdata.MatchedTopic(h.c.Restricted) {
		log.Printf("request [ %s ] was a page [ %s ] with a restricted tag",
			r.URL.Path, h.c.Path+r.URL.Path)
		http.Error(w, "Page not Found", http.StatusNotFound)
		//http.Error(w, err.Error(), http.StatusForbidden)
		return
	}

	// parse any markdown in the input
	body := template.HTML(bodyParseMarkdown(pdata.Page))
	toc := template.HTML(tocParseMarkdown(pdata.Page))
	topics, keywords, authors := pdata.ListMeta()

	// ##TODO## put this template right in the function call
	// Then remove the Page Struct above
	response := Page{
		Title:    pdata.Title,
		ToC:      toc,
		Body:     body,
		Keywords: keywords,
		Topics:   topics,
		Authors:  authors,
	}
	err = allTemplates.ExecuteTemplate(w, h.c.Template, response)
	if err != nil {
		http.Error(w, err.Error(), 500)
	}
}

// RawFile is a http.Handler that serves a raw file back, restricting by file
//  extension if necessary and adding appropiate mime-types.
type RawFile struct {
	c ServerSection
}

func (h RawFile) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// If the request is empty, set it to the default.
	if r.URL.Path == "/" {
		r.URL.Path = path.Clean(h.c.Default)
	}

	for _, restricted := range h.c.Restricted {
		if path.Ext(r.URL.Path) == restricted {
			log.Printf("request %s was has a disallowed extension %s",
				r.URL.Path, restricted)
			http.Error(w, "Request not allowed", 403)
			return
		}
	}

	f, err := os.Open(filepath.Join(h.c.Path, r.URL.Path))
	if err != nil {
		http.Error(w, err.Error(), http.StatusForbidden)
		log.Print(err)
		return
	}

	switch path.Ext(r.URL.Path) {
	case "js":
		w.Header().Set("Content-Type", "text/javascript")
	case "css":
		w.Header().Set("Content-Type", "text/css")
	case "gif":
		w.Header().Set("Content-Type", "image/gif")
	case "png":
		w.Header().Set("Content-Type", "image/png")
	case "jpg":
		w.Header().Set("Content-Type", "image/jpeg")
	case "jpeg":
		w.Header().Set("Content-Type", "image/jpeg")
	}

	_, err = io.Copy(w, f)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		log.Print(err)
	}
}

// FallbackSearchResponse is a function that writes a "bailout" template
func (i *Index) FallbackSearchResponse(w http.ResponseWriter,
	template string) {
	authors, err := i.ListField("author")
	if err != nil {
		http.Error(w, "failed to list authors", http.StatusInternalServerError)
		i.log.Println(err)
		return
	}
	topics, err := i.ListField("topic")
	if err != nil {
		http.Error(w, "failed to list topics", http.StatusInternalServerError)
		i.log.Println(err)
		return
	}

	fields := SearchResponse{Topics: topics, Authors: authors}

	err = allTemplates.ExecuteTemplate(w, template, fields)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		i.log.Println(err)
		return
	}
	return
}
<commit_msg>removed a redudant else statement from the Fields handler<commit_after>package main

import (
	"html/template"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/ajg/form"
)

// Fields is a standard handler that pulls the first folder of the response, and
//  lists that topic or author. If there is none, it falls back to listing all
//  topics or authors with the fallback template.
type Fields struct {
	c ServerSection
	i *Index
}

func (h Fields) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	fields := strings.SplitN(r.URL.Path, "/", 3)

	if len(fields) < 3 || fields[1] == "" {
		// to do if a field was not given
		h.i.FallbackSearchResponse(w, h.c.FallbackTemplate)
		return
	}
	// to be done if a field was given
	results, err := h.i.ListAllField(h.c.Default, fields[0], 100, 1)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	err = allTemplates.ExecuteTemplate(w, h.c.Template, results)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// FuzzySearch is a normal search format - it should provide a point and click
//  interface to allow searching.
type FuzzySearch struct {
	c ServerSection
	i *Index
}

func (h FuzzySearch) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		// to do if a field was not given
		h.i.FallbackSearchResponse(w, h.c.FallbackTemplate)
		return
	}

	err := r.ParseForm()
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	var values FuzzySearchValues
	err = form.DecodeValues(&values, r.Form)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	results, err := h.i.FuzzySearch(values)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	err = allTemplates.ExecuteTemplate(w, h.c.Template, results)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// QuerySearch is a handler that uses a custom search format to do custom queries.
type QuerySearch struct {
	c ServerSection
	i *Index
}

func (h QuerySearch) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	values := struct {
		s        string `form:"s"`
		page     int    `form:"page"`
		pageSize int    `form:"pageSize"`
	}{}

	if r.Method != http.MethodPost {
		// to do if a field was not given
		h.i.FallbackSearchResponse(w, h.c.FallbackTemplate)
		return
	}

	err := r.ParseForm()
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	err = form.DecodeValues(&values, r.Form)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}
	if values.s == "" {
		// to do if a field was not given
		h.i.FallbackSearchResponse(w, h.c.FallbackTemplate)
		return
	}

	results, err := h.i.QuerySearch(values.s, values.page, values.pageSize)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	err = allTemplates.ExecuteTemplate(w, h.c.Template, results)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// Page is a standard data structure used to render markdown pages.
type Page struct {
	Title    string
	ToC      template.HTML
	Body     template.HTML
	Topics   []string
	Keywords []string
	Authors  []string
}

// Markdown is an http.Handler that renders a markdown file and serves it back.
//  Author and Topic tags before the first major title are parsed and displayed.
//  It is possible to restrict access to a page based on topic tag.
type Markdown struct {
	c ServerSection
}

func (h Markdown) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// If the request is empty, set it to the default.
	if r.URL.Path == "/" {
		r.URL.Path = path.Clean(h.c.Default)
	}

	// If the request doesn't end in .md, add that
	if path.Ext(r.URL.Path) != "md" {
		r.URL.Path = r.URL.Path + ".md"
	}

	pdata := new(PageMetadata)
	err := pdata.LoadPage(h.c.Path + r.URL.Path)
	if err != nil {
		log.Printf("request [ %s ] points bad file target [ %s ] sent to server",
			r.URL.Path, h.c.Path)
		http.Error(w, "Page not Found", http.StatusNotFound)
		return
	}

	if pdata.MatchedTopic(h.c.Restricted) {
		log.Printf("request [ %s ] was a page [ %s ] with a restricted tag",
			r.URL.Path, h.c.Path+r.URL.Path)
		http.Error(w, "Page not Found", http.StatusNotFound)
		//http.Error(w, err.Error(), http.StatusForbidden)
		return
	}

	// parse any markdown in the input
	body := template.HTML(bodyParseMarkdown(pdata.Page))
	toc := template.HTML(tocParseMarkdown(pdata.Page))
	topics, keywords, authors := pdata.ListMeta()

	// ##TODO## put this template right in the function call
	// Then remove the Page Struct above
	response := Page{
		Title:    pdata.Title,
		ToC:      toc,
		Body:     body,
		Keywords: keywords,
		Topics:   topics,
		Authors:  authors,
	}
	err = allTemplates.ExecuteTemplate(w, h.c.Template, response)
	if err != nil {
		http.Error(w, err.Error(), 500)
	}
}

// RawFile is a http.Handler that serves a raw file back, restricting by file
//  extension if necessary and adding appropiate mime-types.
type RawFile struct {
	c ServerSection
}

func (h RawFile) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// If the request is empty, set it to the default.
	if r.URL.Path == "/" {
		r.URL.Path = path.Clean(h.c.Default)
	}

	for _, restricted := range h.c.Restricted {
		if path.Ext(r.URL.Path) == restricted {
			log.Printf("request %s was has a disallowed extension %s",
				r.URL.Path, restricted)
			http.Error(w, "Request not allowed", 403)
			return
		}
	}

	f, err := os.Open(filepath.Join(h.c.Path, r.URL.Path))
	if err != nil {
		http.Error(w, err.Error(), http.StatusForbidden)
		log.Print(err)
		return
	}

	switch path.Ext(r.URL.Path) {
	case "js":
		w.Header().Set("Content-Type", "text/javascript")
	case "css":
		w.Header().Set("Content-Type", "text/css")
	case "gif":
		w.Header().Set("Content-Type", "image/gif")
	case "png":
		w.Header().Set("Content-Type", "image/png")
	case "jpg":
		w.Header().Set("Content-Type", "image/jpeg")
	case "jpeg":
		w.Header().Set("Content-Type", "image/jpeg")
	}

	_, err = io.Copy(w, f)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		log.Print(err)
	}
}

// FallbackSearchResponse is a function that writes a "bailout" template
func (i *Index) FallbackSearchResponse(w http.ResponseWriter,
	template string) {
	authors, err := i.ListField("author")
	if err != nil {
		http.Error(w, "failed to list authors", http.StatusInternalServerError)
		i.log.Println(err)
		return
	}
	topics, err := i.ListField("topic")
	if err != nil {
		http.Error(w, "failed to list topics", http.StatusInternalServerError)
		i.log.Println(err)
		return
	}

	fields := SearchResponse{Topics: topics, Authors: authors}

	err = allTemplates.ExecuteTemplate(w, template, fields)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		i.log.Println(err)
		return
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
)

var httpReq HttpReq

type APIHandlerStr struct {
	path string
}

func (af *APIHandlerStr) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	db := connection()
	defer db.Close()
	req, _ := find(db, af.path)
	w.Header().Set("Content-Type", req.ContentType+";charset="+req.CharSet)
	w.WriteHeader(strToInt(req.StatusCode))
	w.Write([]byte(req.Payload))
}

type HandlerFactory struct {
	handler_path string
}

func (hf *HandlerFactory) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == "POST" {
		r.ParseForm()
		httpReq = HttpReq{ApiPath: r.FormValue("api_path"),
			ContentType: r.FormValue("content_type"),
			CharSet:     r.FormValue("char_set"),
			StatusCode:  r.FormValue("status_code"),
			Payload:     r.FormValue("payload"),
		}
		db := connection()
		defer db.Close()
		req, _ := find(db, httpReq.ApiPath)
		if req.ApiPath != "" {
			result := struct {
				Duplicate bool
				Path      string
			}{
				true,
				req.ApiPath,
			}
			data, _ := json.Marshal(result)
			w.Write(data)
			err1 := httpReq.save(db)
			if err1 != nil {
				w.Write([]byte("error!"))
				return
			}
			return
		}
		err1 := httpReq.save(db)
		if err1 != nil {
			w.Write([]byte("error!"))
			return
		}

		hf.handler_path = httpReq.ApiPath
		handler := APIHandlerStr{hf.handler_path}
		handle := fmt.Sprintf("%s", hf.handler_path)
		http.Handle(handle, &handler)
		result := struct {
			Duplicate bool
			Path      string
		}{
			false,
			r.FormValue("api_path"),
		}
		data, _ := json.Marshal(result)
		w.Write(data)
	}
}

func RootHandler(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.NotFound(w, r)
		return
	}

	db := connection()
	defer db.Close()
	a, err := findAll(db)
	if err != nil {
		fmt.Println("ERROR fetching data: ", err)
	}

	data := struct {
		Status  map[int]string
		ApiList []string
	}{
		HttpStatusCodes,
		a,
	}
	t, _ := template.ParseFiles("index.html")
	t.Execute(w, data)
}
<commit_msg>allow CORS<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
)

var httpReq HttpReq

type APIHandlerStr struct {
	path string
}

func (af *APIHandlerStr) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	db := connection()
	defer db.Close()
	req, _ := find(db, af.path)
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Content-Type", req.ContentType+";charset="+req.CharSet)
	w.WriteHeader(strToInt(req.StatusCode))
	w.Write([]byte(req.Payload))
}

type HandlerFactory struct {
	handler_path string
}

func (hf *HandlerFactory) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.Method == "POST" {
		r.ParseForm()
		httpReq = HttpReq{ApiPath: r.FormValue("api_path"),
			ContentType: r.FormValue("content_type"),
			CharSet:     r.FormValue("char_set"),
			StatusCode:  r.FormValue("status_code"),
			Payload:     r.FormValue("payload"),
		}
		db := connection()
		defer db.Close()
		req, _ := find(db, httpReq.ApiPath)
		if req.ApiPath != "" {
			result := struct {
				Duplicate bool
				Path      string
			}{
				true,
				req.ApiPath,
			}
			data, _ := json.Marshal(result)
			w.Write(data)
			err1 := httpReq.save(db)
			if err1 != nil {
				w.Write([]byte("error!"))
				return
			}
			return
		}
		err1 := httpReq.save(db)
		if err1 != nil {
			w.Write([]byte("error!"))
			return
		}

		hf.handler_path = httpReq.ApiPath
		handler := APIHandlerStr{hf.handler_path}
		handle := fmt.Sprintf("%s", hf.handler_path)
		http.Handle(handle, &handler)
		result := struct {
			Duplicate bool
			Path      string
		}{
			false,
			r.FormValue("api_path"),
		}
		data, _ := json.Marshal(result)
		w.Write(data)
	}
}

func RootHandler(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.NotFound(w, r)
		return
	}

	db := connection()
	defer db.Close()
	a, err := findAll(db)
	if err != nil {
		fmt.Println("ERROR fetching data: ", err)
	}

	data := struct {
		Status  map[int]string
		ApiList []string
	}{
		HttpStatusCodes,
		a,
	}
	t, _ := template.ParseFiles("index.html")
	t.Execute(w, data)
}
<|endoftext|>


<commit_before>package logging

import (
	"fmt"
	"io"
	"os"
	"sync"
)

// Handler implementers are sent messages by their owning Logger objects to
// handle however they see fit.  They may ignore the message based on their
// level.
type Handler interface {
	// SetFormatter sets the Formatter to be used for this handler.  Handlers
	// only have one formatter at a time.
	SetFormatter(formatter Formatter)
	// SetLevel sets the logging level that this handler is interested in.
	// Handlers are still given every event that gets to the logger, but they
	// can filter events to a certain level within their Emit methods.
	SetLevel(level Level)
	// Emit is how a Logger feeds its handlers.  Ever event that a logger gets
	// is passed into the Emit method of every Handler.  Handlers must not
	// modify the event because it is shared between the other handlers.
	Emit(event *Event)
}

// HandlerCommon is a struct that contains some common Handler state.
type HandlerCommon struct {
	formatter Formatter
	level     Level
}

// Formatter implements the Handler interface.
func (hc HandlerCommon) Formatter() Formatter {
	return hc.formatter
}

// SetFormatter implements the Handler interface.
func (hc *HandlerCommon) SetFormatter(formatter Formatter) {
	hc.formatter = formatter
}

// Level implements the Handler interface.
func (hc HandlerCommon) Level() Level {
	return hc.level
}

// SetLevel implements the Handler interface.
func (hc *HandlerCommon) SetLevel(level Level) {
	hc.level = level
}

// ConsoleHandler implements the Handler interface by logging events to the
// console.
type ConsoleHandler struct {
	HandlerCommon
}

// Emit implements the Handler interface.
func (ch ConsoleHandler) Emit(event *Event) {
	if event.Level >= ch.level {
		fmt.Fprintf(os.Stderr, ch.formatter.Format(event))
	}
}

// WriterHandler implements the Handler interface by writing events into an
// underlying io.Writer implementation.  Access to the writer is
type WriterHandler struct {
	HandlerCommon

	// L is the mutex that protects the writer from concurrent access.
	L sync.Locker
	w io.Writer
}

type dummyLock struct{}

func (dummyLock) Lock()   {}
func (dummyLock) Unlock() {}

// NewWriterHandler creates a new WriterHandler with an optional lock to
// synchronize access to the writer.  If nil, no locking on the writer is
// performed.
func NewWriterHandler(w io.Writer, lock sync.Locker) *WriterHandler {
	if lock == nil {
		lock = dummyLock{}
	}
	return &WriterHandler{
		HandlerCommon: HandlerCommon{},
		L:             lock,
		w:             w,
	}
}

// Emit implements the Handler interface.
func (wh *WriterHandler) Emit(event *Event) {
	wh.L.Lock()
	defer wh.L.Unlock()
	if event.Level >= wh.level {
		fmt.Fprintf(wh.w, wh.formatter.Format(event))
	}
}
<commit_msg>panic when we fail to emit a log message to an io.Writer<commit_after>package logging

import (
	"fmt"
	"io"
	"os"
	"sync"
)

// Handler implementers are sent messages by their owning Logger objects to
// handle however they see fit.  They may ignore the message based on their
// level.
type Handler interface {
	// SetFormatter sets the Formatter to be used for this handler.  Handlers
	// only have one formatter at a time.
	SetFormatter(formatter Formatter)
	// SetLevel sets the logging level that this handler is interested in.
	// Handlers are still given every event that gets to the logger, but they
	// can filter events to a certain level within their Emit methods.
	SetLevel(level Level)
	// Emit is how a Logger feeds its handlers.  Ever event that a logger gets
	// is passed into the Emit method of every Handler.  Handlers must not
	// modify the event because it is shared between the other handlers.
	Emit(event *Event)
}

// HandlerCommon is a struct that contains some common Handler state.
type HandlerCommon struct {
	formatter Formatter
	level     Level
}

// Formatter implements the Handler interface.
func (hc HandlerCommon) Formatter() Formatter {
	return hc.formatter
}

// SetFormatter implements the Handler interface.
func (hc *HandlerCommon) SetFormatter(formatter Formatter) {
	hc.formatter = formatter
}

// Level implements the Handler interface.
func (hc HandlerCommon) Level() Level {
	return hc.level
}

// SetLevel implements the Handler interface.
func (hc *HandlerCommon) SetLevel(level Level) {
	hc.level = level
}

// ConsoleHandler implements the Handler interface by logging events to the
// console.
type ConsoleHandler struct {
	HandlerCommon
}

// Emit implements the Handler interface.
func (ch ConsoleHandler) Emit(event *Event) {
	if event.Level >= ch.level {
		fmt.Fprintf(os.Stderr, ch.formatter.Format(event))
	}
}

// WriterHandler implements the Handler interface by writing events into an
// underlying io.Writer implementation.  Access to the writer is
type WriterHandler struct {
	HandlerCommon

	// L is the mutex that protects the writer from concurrent access.
	L sync.Locker
	w io.Writer
}

type dummyLock struct{}

func (dummyLock) Lock()   {}
func (dummyLock) Unlock() {}

// NewWriterHandler creates a new WriterHandler with an optional lock to
// synchronize access to the writer.  If nil, no locking on the writer is
// performed.
func NewWriterHandler(w io.Writer, lock sync.Locker) *WriterHandler {
	if lock == nil {
		lock = dummyLock{}
	}
	return &WriterHandler{
		HandlerCommon: HandlerCommon{},
		L:             lock,
		w:             w,
	}
}

// Emit implements the Handler interface.
func (wh *WriterHandler) Emit(event *Event) {
	wh.L.Lock()
	defer wh.L.Unlock()
	if event.Level >= wh.level {
		_, err := fmt.Fprintf(wh.w, wh.formatter.Format(event))
		if err != nil {
			panic(err)
		}
	}
}
<|endoftext|>


<commit_before>package dochaincore

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"
	"sync"
)

const indexPageHTML = `<!DOCTYPE html>

		<title>One-Click Chain Core DigitalOcean</title>
		<link rel="stylesheet" href="/static/style.css">

Install Chain Core on a DigitalOcean droplet. This installer creates a new 1gb droplet and a 100gb block storage volume on your DigitalOcean account. It installs Chain Core on the droplet using the attached volume for storage. The approximate cost on DigitalOcean is $20/month.

			Install Chain Core

`
const progressPageHTML = `
<!DOCTYPE html>

		<title>One-Click Chain Core DigitalOcean</title>
		<link rel="stylesheet" href="/static/style.css">

Initializing droplet…

Success! Chain Core has been installed on your DigitalOcean droplet. To access
				Chain Core's API and Dashboard, you'll need your client token:

				Open dashboard

`

func Handler(oauthClientID, oauthClientSecret, host string) http.Handler {
	h := &handler{
		oauthClientID:     oauthClientID,
		oauthClientSecret: oauthClientSecret,
		host:              host,
		progressTmpl:      template.Must(template.New("progresspage").Parse(progressPageHTML)),
		indexTmpl:         template.Must(template.New("index").Parse(indexPageHTML)),
		installs:          make(map[string]*install),
	}
	mux := http.NewServeMux()
	mux.HandleFunc("/status/", h.status)
	mux.HandleFunc("/grant", h.grant)
	mux.HandleFunc("/progress/", h.progressPage)
	mux.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.Dir("static"))))
	mux.HandleFunc("/", h.index)
	return mux
}

type handler struct {
	oauthClientID     string
	oauthClientSecret string
	host              string
	progressTmpl      *template.Template
	indexTmpl         *template.Template

	installMu sync.Mutex
	installs  map[string]*install
}

func (h *handler) index(rw http.ResponseWriter, req *http.Request) {
	b := make([]byte, 16)
	_, err := rand.Read(b)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}
	installID := hex.EncodeToString(b)

	h.installMu.Lock()
	h.installs[installID] = &install{Status: "pending auth"}
	h.installMu.Unlock()

	vals := make(url.Values)
	vals.Set("response_type", "code")
	vals.Set("client_id", h.oauthClientID)
	vals.Set("state", installID)
	vals.Set("scope", "read write")
	vals.Set("redirect_uri", h.host+"/progress")
	u := url.URL{
		Scheme:   "https",
		Host:     "cloud.digitalocean.com",
		Path:     "/v1/oauth/authorize",
		RawQuery: vals.Encode(),
	}

	tmplData := struct {
		InstallLink string
	}{
		InstallLink: u.String(),
	}
	h.indexTmpl.Execute(rw, tmplData)
}

func (h *handler) grant(rw http.ResponseWriter, req *http.Request) {
	code, state := req.FormValue("code"), req.FormValue("state")
	if code == "" || state == "" {
		http.Error(rw, "invalid oauth2 grant", http.StatusBadRequest)
		return
	}

	h.installMu.Lock()
	curr := h.installs[state]
	h.installMu.Unlock()
	if curr == nil {
		http.Error(rw, "invalid oauth2 state", http.StatusBadRequest)
		return
	}

	// Claim the code grant
	vals := make(url.Values)
	vals.Set("grant_type", "authorization_code")
	vals.Set("code", code)
	vals.Set("client_id", h.oauthClientID)
	vals.Set("client_secret", h.oauthClientSecret)
	vals.Set("redirect_uri", h.host+"/progress")
	u := url.URL{
		Scheme:   "https",
		Host:     "cloud.digitalocean.com",
		Path:     "/v1/oauth/token",
		RawQuery: vals.Encode(),
	}
	resp, err := http.Post(u.String(), "application/x-www-form-urlencoded", nil)
	if err != nil {
		http.Error(rw, "internal server error", http.StatusInternalServerError)
		return
	}
	defer resp.Body.Close()

	var decodedResponse struct {
		AccessToken  string `json:"access_token"`
		TokenType    string `json:"bearer"`
		ExpiresIn    int    `json:"expires_in"`
		RefreshToken string `json:"refresh_token"`
		Scope        string `json:"scope"`
		Info         struct {
			Name  string `json:"name"`
			Email string `json:"email"`
			UUID  string `json:"uuid"`
		} `json:"info"`
	}
	err = json.NewDecoder(resp.Body).Decode(&decodedResponse)
	if err != nil {
		http.Error(rw, "err decoding access token grant", http.StatusBadRequest)
		return
	}
	if decodedResponse.AccessToken == "" {
		http.Error(rw, "missing access token", http.StatusBadRequest)
		return
	}
	if decodedResponse.Scope != "read write" {
		http.Error(rw, "need read write OAuth scope", http.StatusBadRequest)
		return
	}

	curr.mu.Lock()
	curr.accessToken = decodedResponse.AccessToken
	curr.mu.Unlock()

	http.Redirect(rw, req, "/install/"+state, http.StatusFound)
}

func (h *handler) progressPage(rw http.ResponseWriter, req *http.Request) {
	id := path.Base(req.URL.Path)
	h.installMu.Lock()
	curr := h.installs[id]
	h.installMu.Unlock()

	if curr == nil {
		http.NotFound(rw, req)
		return
	}

	go curr.init(id)

	tmplData := struct {
		InstallID string
	}{
		InstallID: id,
	}
	err := h.progressTmpl.Execute(rw, tmplData)
	if err != nil {
		fmt.Fprintf(os.Stderr, "executing template: %s", err.Error())
	}
}

func (h *handler) status(rw http.ResponseWriter, req *http.Request) {
	id := path.Base(req.URL.Path)
	h.installMu.Lock()
	curr := h.installs[id]
	h.installMu.Unlock()

	if curr == nil {
		http.NotFound(rw, req)
		return
	}

	// Marshal to a buffer first so that a really slow request can't
	// keep curr.mu locked indefinitely.
	var buf bytes.Buffer
	curr.mu.Lock()
	_ = json.NewEncoder(&buf).Encode(curr)
	curr.mu.Unlock()

	rw.Header().Set("Content-Type", "application/json")
	rw.Write(buf.Bytes())
}

type install struct {
	mu          sync.Mutex
	Status      string `json:"status"`
	ClientToken string `json:"client_token"`
	IPAddress   string `json:"ip_address"`
	accessToken string
	c           *Core
}

func (i *install) setStatus(status string) {
	i.mu.Lock()
	defer i.mu.Unlock()
	i.Status = status
}

func (i *install) init(state string) {
	defer revoke(i.accessToken)

	var core *Core
	var err error

	defer func() {
		if err != nil {
			i.setStatus(err.Error())
		}
	}()

	// Start deploying and create the droplet.
	core, err = Deploy(i.accessToken, DropletName("chain-core-"+state[:6]))
	if err != nil {
		return
	}

	i.mu.Lock()
	i.IPAddress = core.IPv4Address
	i.c = core
	i.Status = "waiting for ssh"
	i.mu.Unlock()

	err = WaitForSSH(core)
	if err != nil {
		return
	}

	i.setStatus("waiting for http")
	err = WaitForHTTP(core)
	if err != nil {
		return
	}

	i.setStatus("creating client token")
	token, err := CreateClientToken(core)
	if err != nil {
		return
	}

	i.mu.Lock()
	i.Status = "done"
	i.ClientToken = token
	i.c = nil // garbage collect the SSH keys
	i.mu.Unlock()
}

func revoke(accessToken string) error {
	body := strings.NewReader(url.Values{"token": {accessToken}}.Encode())
	req, err := http.NewRequest("POST", "https://cloud.digitalocean.com/v1/oauth/revoke", body)
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", "Bearer "+accessToken)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("revoke endpoint returned %d status code", resp.StatusCode)
	}
	return nil
}
<commit_msg>fix redirect URI<commit_after>package dochaincore

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"html/template"
	"net/http"
	"net/url"
	"os"
	"path"
	"strings"
	"sync"
)

const indexPageHTML = `<!DOCTYPE html>

		<title>One-Click Chain Core DigitalOcean</title>
		<link rel="stylesheet" href="/static/style.css">

Install Chain Core on a DigitalOcean droplet. This installer creates a new 1gb droplet and a 100gb block storage volume on your DigitalOcean account. It installs Chain Core on the droplet using the attached volume for storage. The approximate cost on DigitalOcean is $20/month.

			Install Chain Core

`
const progressPageHTML = `
<!DOCTYPE html>

		<title>One-Click Chain Core DigitalOcean</title>
		<link rel="stylesheet" href="/static/style.css">

Initializing droplet…

Success! Chain Core has been installed on your DigitalOcean droplet. To access
				Chain Core's API and Dashboard, you'll need your client token:

				Open dashboard

`

func Handler(oauthClientID, oauthClientSecret, host string) http.Handler {
	h := &handler{
		oauthClientID:     oauthClientID,
		oauthClientSecret: oauthClientSecret,
		host:              host,
		progressTmpl:      template.Must(template.New("progresspage").Parse(progressPageHTML)),
		indexTmpl:         template.Must(template.New("index").Parse(indexPageHTML)),
		installs:          make(map[string]*install),
	}
	mux := http.NewServeMux()
	mux.HandleFunc("/status/", h.status)
	mux.HandleFunc("/grant", h.grant)
	mux.HandleFunc("/progress/", h.progressPage)
	mux.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.Dir("static"))))
	mux.HandleFunc("/", h.index)
	return mux
}

type handler struct {
	oauthClientID     string
	oauthClientSecret string
	host              string
	progressTmpl      *template.Template
	indexTmpl         *template.Template

	installMu sync.Mutex
	installs  map[string]*install
}

func (h *handler) index(rw http.ResponseWriter, req *http.Request) {
	b := make([]byte, 16)
	_, err := rand.Read(b)
	if err != nil {
		http.Error(rw, err.Error(), http.StatusInternalServerError)
		return
	}
	installID := hex.EncodeToString(b)

	h.installMu.Lock()
	h.installs[installID] = &install{Status: "pending auth"}
	h.installMu.Unlock()

	vals := make(url.Values)
	vals.Set("response_type", "code")
	vals.Set("client_id", h.oauthClientID)
	vals.Set("state", installID)
	vals.Set("scope", "read write")
	vals.Set("redirect_uri", h.host+"/grant")
	u := url.URL{
		Scheme:   "https",
		Host:     "cloud.digitalocean.com",
		Path:     "/v1/oauth/authorize",
		RawQuery: vals.Encode(),
	}

	tmplData := struct {
		InstallLink string
	}{
		InstallLink: u.String(),
	}
	h.indexTmpl.Execute(rw, tmplData)
}

func (h *handler) grant(rw http.ResponseWriter, req *http.Request) {
	code, state := req.FormValue("code"), req.FormValue("state")
	if code == "" || state == "" {
		http.Error(rw, "invalid oauth2 grant", http.StatusBadRequest)
		return
	}

	h.installMu.Lock()
	curr := h.installs[state]
	h.installMu.Unlock()
	if curr == nil {
		http.Error(rw, "invalid oauth2 state", http.StatusBadRequest)
		return
	}

	// Claim the code grant
	vals := make(url.Values)
	vals.Set("grant_type", "authorization_code")
	vals.Set("code", code)
	vals.Set("client_id", h.oauthClientID)
	vals.Set("client_secret", h.oauthClientSecret)
	vals.Set("redirect_uri", h.host+"/grant")
	u := url.URL{
		Scheme:   "https",
		Host:     "cloud.digitalocean.com",
		Path:     "/v1/oauth/token",
		RawQuery: vals.Encode(),
	}
	resp, err := http.Post(u.String(), "application/x-www-form-urlencoded", nil)
	if err != nil {
		http.Error(rw, "internal server error", http.StatusInternalServerError)
		return
	}
	defer resp.Body.Close()

	var decodedResponse struct {
		AccessToken  string `json:"access_token"`
		TokenType    string `json:"bearer"`
		ExpiresIn    int    `json:"expires_in"`
		RefreshToken string `json:"refresh_token"`
		Scope        string `json:"scope"`
		Info         struct {
			Name  string `json:"name"`
			Email string `json:"email"`
			UUID  string `json:"uuid"`
		} `json:"info"`
	}
	err = json.NewDecoder(resp.Body).Decode(&decodedResponse)
	if err != nil {
		http.Error(rw, "err decoding access token grant", http.StatusBadRequest)
		return
	}
	if decodedResponse.AccessToken == "" {
		http.Error(rw, "missing access token", http.StatusBadRequest)
		return
	}
	if decodedResponse.Scope != "read write" {
		http.Error(rw, "need read write OAuth scope", http.StatusBadRequest)
		return
	}

	curr.mu.Lock()
	curr.accessToken = decodedResponse.AccessToken
	curr.mu.Unlock()

	http.Redirect(rw, req, "/install/"+state, http.StatusFound)
}

func (h *handler) progressPage(rw http.ResponseWriter, req *http.Request) {
	id := path.Base(req.URL.Path)
	h.installMu.Lock()
	curr := h.installs[id]
	h.installMu.Unlock()

	if curr == nil {
		http.NotFound(rw, req)
		return
	}

	go curr.init(id)

	tmplData := struct {
		InstallID string
	}{
		InstallID: id,
	}
	err := h.progressTmpl.Execute(rw, tmplData)
	if err != nil {
		fmt.Fprintf(os.Stderr, "executing template: %s", err.Error())
	}
}

func (h *handler) status(rw http.ResponseWriter, req *http.Request) {
	id := path.Base(req.URL.Path)
	h.installMu.Lock()
	curr := h.installs[id]
	h.installMu.Unlock()

	if curr == nil {
		http.NotFound(rw, req)
		return
	}

	// Marshal to a buffer first so that a really slow request can't
	// keep curr.mu locked indefinitely.
	var buf bytes.Buffer
	curr.mu.Lock()
	_ = json.NewEncoder(&buf).Encode(curr)
	curr.mu.Unlock()

	rw.Header().Set("Content-Type", "application/json")
	rw.Write(buf.Bytes())
}

type install struct {
	mu          sync.Mutex
	Status      string `json:"status"`
	ClientToken string `json:"client_token"`
	IPAddress   string `json:"ip_address"`
	accessToken string
	c           *Core
}

func (i *install) setStatus(status string) {
	i.mu.Lock()
	defer i.mu.Unlock()
	i.Status = status
}

func (i *install) init(state string) {
	defer revoke(i.accessToken)

	var core *Core
	var err error

	defer func() {
		if err != nil {
			i.setStatus(err.Error())
		}
	}()

	// Start deploying and create the droplet.
	core, err = Deploy(i.accessToken, DropletName("chain-core-"+state[:6]))
	if err != nil {
		return
	}

	i.mu.Lock()
	i.IPAddress = core.IPv4Address
	i.c = core
	i.Status = "waiting for ssh"
	i.mu.Unlock()

	err = WaitForSSH(core)
	if err != nil {
		return
	}

	i.setStatus("waiting for http")
	err = WaitForHTTP(core)
	if err != nil {
		return
	}

	i.setStatus("creating client token")
	token, err := CreateClientToken(core)
	if err != nil {
		return
	}

	i.mu.Lock()
	i.Status = "done"
	i.ClientToken = token
	i.c = nil // garbage collect the SSH keys
	i.mu.Unlock()
}

func revoke(accessToken string) error {
	body := strings.NewReader(url.Values{"token": {accessToken}}.Encode())
	req, err := http.NewRequest("POST", "https://cloud.digitalocean.com/v1/oauth/revoke", body)
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", "Bearer "+accessToken)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("revoke endpoint returned %d status code", resp.StatusCode)
	}
	return nil
}
<|endoftext|>


<commit_before>package dhcpv4

import "net"

// PacketReader defines the ReadFrom function as defined in net.PacketConn.
type PacketReader interface {
	ReadFrom(b []byte) (n int, addr net.Addr, err error)
}

// PacketWriter defines the ReadFrom function as defined in net.PacketConn.
type PacketWriter interface {
	WriteTo(b []byte, addr net.Addr) (n int, err error)
}

// PacketConn groups PacketReader and PacketWriter to form a subset of net.PacketConn.
type PacketConn interface {
	PacketReader
	PacketWriter
}

type replyWriter struct {
	pw PacketWriter

	// The client address, if any
	addr net.UDPAddr
}

func (rw *replyWriter) WriteReply(r Reply) error {
	var err error

	err = r.Validate()
	if err != nil {
		return err
	}

	bytes, err := r.ToBytes()
	if err != nil {
		return err
	}

	req := r.Request()
	addr := rw.addr
	bcast := req.Flags()[0] & 128

	// Broadcast the reply if the request packet has no address associated with
	// it, or if the client explicitly asks for a broadcast reply.
	if addr.IP.Equal(net.IPv4zero) || bcast > 0 {
		addr.IP = net.IPv4bcast
	}

	_, err = rw.pw.WriteTo(bytes, &addr)
	if err != nil {
		return err
	}

	return nil
}

// Handler defines the interface an object needs to implement to handle DHCP
// packets. The handler should do a type switch on the Request object that is
// passed as argument to determine what kind of packet it is dealing with. It
// can use the WriteReply function on the request to send a reply back to the
// peer responsible for sending the request packet. While the handler may be
// blocking, it is not encouraged. Rather, the handler should return as soon as
// possible to avoid blocking the serve loop. If blocking operations need to be
// executed to determine if the request packet needs a reply, and if so, what
// kind of reply, it is recommended to handle this in separate goroutines. The
// WriteReply function can be called from multiple goroutines without needing
// extra synchronization.
type Handler interface {
	ServeDHCP(req Request)
}

// Serve reads packets off the network and passes them to the specified
// handler. It is up to the handler to packets to per-client serve loops, if
// that is what you want.
func Serve(pc PacketConn, h Handler) error {
	buf := make([]byte, 65536)

	for {
		n, addr, err := pc.ReadFrom(buf)
		if err != nil {
			return err
		}

		p, err := PacketFromBytes(buf[:n])
		if err != nil {
			continue
		}

		// Filter everything but requests
		if OpCode(p.Op()[0]) != BootRequest {
			continue
		}

		rw := replyWriter{
			pw:   pc,
			addr: *addr.(*net.UDPAddr),
		}

		var req Request

		switch p.GetMessageType() {
		case MessageTypeDHCPDiscover:
			req = DHCPDiscover{p, &rw}
		case MessageTypeDHCPRequest:
			req = DHCPRequest{p, &rw}
		case MessageTypeDHCPDecline:
			req = DHCPDecline{p}
		case MessageTypeDHCPRelease:
			req = DHCPRelease{p}
		case MessageTypeDHCPInform:
			req = DHCPInform{p, &rw}
		}

		if req != nil {
			h.ServeDHCP(req)
		}
	}
}
<commit_msg>Add PacketConn filter to filter packets by interface<commit_after>package dhcpv4

import (
	"fmt"
	"net"

	"code.google.com/p/go.net/ipv4"
)

// PacketReader defines the ReadFrom function as defined in net.PacketConn.
type PacketReader interface {
	ReadFrom(b []byte) (n int, addr net.Addr, err error)
}

// PacketWriter defines the ReadFrom function as defined in net.PacketConn.
type PacketWriter interface {
	WriteTo(b []byte, addr net.Addr) (n int, err error)
}

// PacketConn groups PacketReader and PacketWriter to form a subset of net.PacketConn.
type PacketConn interface {
	PacketReader
	PacketWriter
}

type replyWriter struct {
	pw PacketWriter

	// The client address, if any
	addr net.UDPAddr
}

func (rw *replyWriter) WriteReply(r Reply) error {
	var err error

	err = r.Validate()
	if err != nil {
		return err
	}

	bytes, err := r.ToBytes()
	if err != nil {
		return err
	}

	req := r.Request()
	addr := rw.addr
	bcast := req.Flags()[0] & 128

	// Broadcast the reply if the request packet has no address associated with
	// it, or if the client explicitly asks for a broadcast reply.
	if addr.IP.Equal(net.IPv4zero) || bcast > 0 {
		addr.IP = net.IPv4bcast
	}

	_, err = rw.pw.WriteTo(bytes, &addr)
	if err != nil {
		return err
	}

	return nil
}

// Handler defines the interface an object needs to implement to handle DHCP
// packets. The handler should do a type switch on the Request object that is
// passed as argument to determine what kind of packet it is dealing with. It
// can use the WriteReply function on the request to send a reply back to the
// peer responsible for sending the request packet. While the handler may be
// blocking, it is not encouraged. Rather, the handler should return as soon as
// possible to avoid blocking the serve loop. If blocking operations need to be
// executed to determine if the request packet needs a reply, and if so, what
// kind of reply, it is recommended to handle this in separate goroutines. The
// WriteReply function can be called from multiple goroutines without needing
// extra synchronization.
type Handler interface {
	ServeDHCP(req Request)
}

// Serve reads packets off the network and calls the specified handler.
func Serve(pc PacketConn, h Handler) error {
	buf := make([]byte, 65536)

	for {
		n, addr, err := pc.ReadFrom(buf)
		if err != nil {
			return err
		}

		p, err := PacketFromBytes(buf[:n])
		if err != nil {
			continue
		}

		// Filter everything but requests
		if OpCode(p.Op()[0]) != BootRequest {
			continue
		}

		rw := replyWriter{
			pw:   pc,
			addr: *addr.(*net.UDPAddr),
		}

		var req Request

		switch p.GetMessageType() {
		case MessageTypeDHCPDiscover:
			req = DHCPDiscover{p, &rw}
		case MessageTypeDHCPRequest:
			req = DHCPRequest{p, &rw}
		case MessageTypeDHCPDecline:
			req = DHCPDecline{p}
		case MessageTypeDHCPRelease:
			req = DHCPRelease{p}
		case MessageTypeDHCPInform:
			req = DHCPInform{p, &rw}
		}

		if req != nil {
			h.ServeDHCP(req)
		}
	}
}

// packetConnFilter wraps net.PacketConn and only reads and writes packet from
// and to the specified network interface.
type packetConnFilter struct {
	net.PacketConn

	ipv4pc *ipv4.PacketConn
	ipv4cm *ipv4.ControlMessage
}

// ReadFrom reads a packet from the connection copying the payload into b. It
// inherits its semantics from ipv4.PacketConn and subsequently net.PacketConn,
// but filters out packets that arrived on an interface other than the one
// specified in the packetConnFilter structure.
func (p *packetConnFilter) ReadFrom(b []byte) (n int, addr net.Addr, err error) {
	for {
		n, cm, src, err := p.ipv4pc.ReadFrom(b)
		if err != nil {
			return n, src, err
		}

		// Read another packet if it didn't arrive on the right interface
		if cm.IfIndex != p.ipv4cm.IfIndex {
			continue
		}

		return n, src, err
	}
}

// WriteTo writes a packet with payload b to addr. It inherits its semantics
// from ipv4.PacketConn and subsequently net.PacketConn, but explicitly sends
// the packet over the interface specified in the packetConnFilter structure.
func (p *packetConnFilter) WriteTo(b []byte, addr net.Addr) (n int, err error) {
	return p.ipv4pc.WriteTo(b, p.ipv4cm, addr)
}

// PacketConnFilter wraps a net.PacketConn and only reads packets from and
// writes packets to the network interface associated with the specified IP
// address. It may return an error if it cannot initialize the underlying
// socket correctly. It panics if it cannot find the network interface
// associated with the specified IP.
func PacketConnFilter(pc net.PacketConn, ip net.IP) (net.PacketConn, error) {
	ipv4pc := ipv4.NewPacketConn(pc)
	if err := ipv4pc.SetControlMessage(ipv4.FlagInterface, true); err != nil {
		return nil, err
	}

	p := packetConnFilter{
		PacketConn: pc,

		ipv4pc: ipv4pc,
		ipv4cm: &ipv4.ControlMessage{
			IfIndex: LookupInterfaceIndexForIP(ip),
		},
	}

	return &p, nil
}

// LookupInterfaceIndexForIP finds the system-wide network interface index that
// is associated with the specified IP address.
func LookupInterfaceIndexForIP(ip net.IP) int {
	is, err := net.Interfaces()
	if err != nil {
		panic(err)
	}
	for _, i := range is {
		as, err := i.Addrs()
		if err != nil {
			panic(err)
		}
		for _, a := range as {
			if a.(*net.IPNet).IP.String() == ip.String() {
				return i.Index
			}
		}
	}

	// Not really a recoverable error...
	panic(fmt.Sprintf("dhcpv4: can't find network interface for: %s", ip))
}
<|endoftext|>