<commit_before>package main

import (
  "log"
  "os"
  "fmt"
  "time"
  "strings"
  "net/http"
  "io/ioutil"
  "regexp"
  "encoding/json"
  "gopkg.in/mgo.v2"
  "gopkg.in/mgo.v2/bson"
  "github.com/gorilla/mux"
)

/**
 *  Shuttle Tracker
 *   Auto Updater - send request to iTrak API,
 *                  get updated shuttle info,
 *                  store updated records in db
 */
type VehicleUpdate struct {
  Id          bson.ObjectId                 `bson:"_id,omitempty"`
  VehicleId   string     `json:"vehicleId"   bson:"vehicleId,omitempty"`
  Lat         string     `json:"lat"         bson:"lat"`
  Lng         string     `json:"lng"         bson:"lng"`
  Heading     string     `json:"heading"     bson:"heading"`
  Speed       string     `json:"speed"       bson:"speed"`
  Lock        string     `json:"lock"        bson:"lock"`
  Time        string     `json:"time"        bson:"time"`
  Date        string     `json:"date"        bson:"date"`
  Status      string     `json:"status"      bson:"status"`
  Created     time.Time  `json:"created"     bson:"created"`
}

func (App *App) UpdateShuttles(dataFeed string, updateInterval int) {
  var st time.Duration
  for {

    // Sleep for n seconds before updating again
    time.Sleep(st)
    if st == 0 {
      // Initialize the sleep timer after the first sleep.  This lets us sleep during errors
      // when we 'continue' back to the top of the loop without waiting to sleep for the first
      // update run.
      st = time.Duration(updateInterval) * time.Second
    }

    // Make request to our tracking data feed
    resp, err := http.Get(dataFeed)
    if err != nil {
      log.Printf("error getting data feed: %v", err)
      continue;
    }
    defer resp.Body.Close()

    // Read response body content
    body, err := ioutil.ReadAll(resp.Body)
    if err != nil {
      log.Printf("error reading data feed: %v", err)
      continue;
    }

    delim := "eof"
    // Iterate through all vehicles returned by data feed
    vehicles_data := strings.Split(string(body), delim)
    for i := 1; i < len(vehicles_data)-1; i++ {

      // Match eatch API field with any number (+)
      //   of the previous expressions (\d digit, \. escaped period, - negative number)
      //   Specify named capturing groups to store each field from data feed
      re := regexp.MustCompile(`(?P<id>Vehicle ID:([\d\.]+)) (?P<lat>lat:([\d\.-]+)) (?P<lng>lon:([\d\.-]+)) (?P<heading>dir:([\d\.-]+)) (?P<speed>spd:([\d\.-]+)) (?P<lock>lck:([\d\.-]+)) (?P<time>time:([\d]+)) (?P<date>date:([\d]+)) (?P<status>trig:([\d]+))`)
      n := re.SubexpNames()
      match := re.FindAllStringSubmatch(vehicles_data[i], -1)[0]

      // Store named capturing group and matching expression as a key value pair
      result := map[string]string{}
      for i, item := range match {
        result[n[i]] = item
      }

      // Create new vehicle update & insert update into database
      update := VehicleUpdate {
        Id:        bson.NewObjectId(),
        VehicleId: strings.Replace(result["id"], "Vehicle ID:", "", -1),
        Lat:       strings.Replace(result["lat"], "lat:", "", -1),
        Lng:       strings.Replace(result["lng"], "lon:", "", -1),
        Heading:   strings.Replace(result["heading"], "dir:", "", -1),
        Speed:     strings.Replace(result["speed"], "spd:", "", -1),
        Lock:      strings.Replace(result["lock"], "lck:", "", -1),
        Time:      strings.Replace(result["time"], "time:", "", -1),
        Date:      strings.Replace(result["date"], "date:", "", -1),
        Status:    strings.Replace(result["status"], "trig:", "", -1),
        Created:   time.Now()}

      err := App.Updates.Insert(&update)

      if err != nil {
        fmt.Println(err.Error())
      }
    }
  }
}

/**
 *  Route handlers - API requests,
 *                   serve view files
 */

type Vehicle struct {
  Id          bson.ObjectId                 `bson:"_id,omitempty"`
  VehicleId   string     `json:"vehicleId"   bson:"vehicleId,omitempty"`
  VehicleName string     `json:"vehicleName" bson:"vehicleName"`
  Created     time.Time  `json:"created"     bson:"created"`
}

/**
 * Find all vehicles in the database
 *
 */
func (App *App) VehiclesHandler(w http.ResponseWriter, r *http.Request) {
  // Find all vehicles in database
  var vehicles []Vehicle
  err := App.Vehicles.Find(bson.M{}).All(&vehicles)
  // Handle query errors
  if err != nil {
    http.Error(w, err.Error(), http.StatusInternalServerError)
  }
  // Send each vehicle to client as JSON
  vehiclesJSON, err := json.MarshalIndent(vehicles, "", " ")
  fmt.Fprintf(w, string(vehiclesJSON))
}

/**
 * Add a new vehicle to the database
 *
 */
func (App *App) VehiclesCreateHandler(w http.ResponseWriter, r *http.Request) {
  // Create new vehicle object using request fields
  vehicle := Vehicle{}
  vehicleData := json.NewDecoder(r.Body)
  err := vehicleData.Decode(&vehicle)
  // Error handling
  if err != nil {
    http.Error(w, err.Error(), http.StatusInternalServerError)
  }
  // Store new vehicle under vehicles collection
  err = App.Vehicles.Insert(&vehicle)
  // Error handling
  if err != nil {
    http.Error(w, err.Error(), http.StatusInternalServerError)
  }
}

/**
 * Vehicle Updates - Get most recent update for each
 *                   vehicle in the vehicles collection
 */
func (App *App) UpdatesHandler(w http.ResponseWriter, r *http.Request) {
  // Store updates for each vehicle
  var vehicles []Vehicle
  var updates []VehicleUpdate
  var update VehicleUpdate
  // Query all Vehicles
  err := App.Vehicles.Find(bson.M{}).All(&vehicles)
  // Handle errors
  if err != nil {
    http.Error(w, err.Error(), http.StatusInternalServerError)
  }
  // Find recent updates for each vehicle
  for _,vehicle := range vehicles {
    err := App.Updates.Find(bson.M{"vehicleId": vehicle.VehicleId}).Sort("-created").Limit(1).One(&update)

    if err == nil {
      updates = append(updates, update)
    }
  }
  // Convert updates to JSON
  u, err := json.MarshalIndent(updates, "", " ")
  // Handle JSON parsing errors
  if err != nil {
    http.Error(w, err.Error(), http.StatusInternalServerError)
  }
  // Send updates to client
  fmt.Fprint(w, string(u))
}

func IndexHandler(w http.ResponseWriter, r *http.Request) {
  http.ServeFile(w, r, "index.html")
}

func AdminHandler(w http.ResponseWriter, r *http.Request) {
  http.ServeFile(w, r, "admin.html")
}

/**
 *  Main - connect to database,
 *         handle routing,
 *         start tracker go routine,
 *         serve requests
 */

type Configuration struct {
  DataFeed        string
  UpdateInterval  int
  MongoUrl        string
  MongoPort       string
}

type App struct {
  Session    *mgo.Session
  Updates    *mgo.Collection
  Vehicles   *mgo.Collection
}

func ReadConfiguration(fileName string) (*Configuration, error) {
  // Open config file and decode JSON to Configuration struct
  file, err := os.Open(fileName)
  if err != nil {
    return nil, err
  }
  decoder := json.NewDecoder(file)
  config := Configuration{}
  if err := decoder.Decode(&config); err != nil {
    return nil, err
  }
  return &config, nil
}

func main() {
  // Read app configuration file
  config, err := ReadConfiguration("conf.json")
  if err != nil {
    log.Fatalf("error reading configuration file: %v", err)
  }

  // Connect to MongoDB
  session, err := mgo.Dial(config.MongoUrl + ":" + config.MongoPort)
  if err != nil {
    log.Fatalf("mongoDB connection failed: %v", err)
  }
  // close Mongo session when server terminates
  defer session.Close()

  // Create Shuttles object to store database session and collections
  App := &App{
    session,
    session.DB("shuttle_tracking").C("updates"),
    session.DB("shuttle_tracking").C("vehicles"),
  }

  // Start auto updater
  go App.UpdateShuttles(config.DataFeed, config.UpdateInterval)

  // Routing
  r := mux.NewRouter()
  r.HandleFunc("/", IndexHandler).Methods("GET")
  r.HandleFunc("/admin", AdminHandler).Methods("GET")
  r.HandleFunc("/admin/vehicles", AdminHandler).Methods("GET")
  r.HandleFunc("/admin/tracking", AdminHandler).Methods("GET")
  r.HandleFunc("/admin/stops", AdminHandler).Methods("GET")
  r.HandleFunc("/vehicles", App.VehiclesHandler).Methods("GET")
  r.HandleFunc("/vehicles/create", App.VehiclesCreateHandler).Methods("POST")
  r.HandleFunc("/updates", App.UpdatesHandler).Methods("GET")
  // Static files
  r.PathPrefix("/bower_components/").Handler(http.StripPrefix("/bower_components/", http.FileServer(http.Dir("bower_components/"))))
  r.PathPrefix("/static/").Handler(http.StripPrefix("/static/", http.FileServer(http.Dir("static/"))))
  // Serve requests
  http.Handle("/", r)
  if err := http.ListenAndServe(":8080", r); err != nil {
    log.Fatalf("Unable to ListenAndServe: %v", err)
  }
}
<commit_msg>Improve vehicle update effeciency and logging.<commit_after>package main

import (
  "log"
  "os"
  "fmt"
  "time"
  "strings"
  "net/http"
  "io/ioutil"
  "regexp"
  "encoding/json"
  "gopkg.in/mgo.v2"
  "gopkg.in/mgo.v2/bson"
  "github.com/gorilla/mux"
)

/**
 *  Shuttle Tracker
 *   Auto Updater - send request to iTrak API,
 *                  get updated shuttle info,
 *                  store updated records in db
 */
type VehicleUpdate struct {
  Id          bson.ObjectId                 `bson:"_id,omitempty"`
  VehicleId   string     `json:"vehicleId"   bson:"vehicleId,omitempty"`
  Lat         string     `json:"lat"         bson:"lat"`
  Lng         string     `json:"lng"         bson:"lng"`
  Heading     string     `json:"heading"     bson:"heading"`
  Speed       string     `json:"speed"       bson:"speed"`
  Lock        string     `json:"lock"        bson:"lock"`
  Time        string     `json:"time"        bson:"time"`
  Date        string     `json:"date"        bson:"date"`
  Status      string     `json:"status"      bson:"status"`
  Created     time.Time  `json:"created"     bson:"created"`
}

var (
  // Match each API field with any number (+)
  //   of the previous expressions (\d digit, \. escaped period, - negative number)
  //   Specify named capturing groups to store each field from data feed
  dataRe = regexp.MustCompile(`(?P<id>Vehicle ID:([\d\.]+)) (?P<lat>lat:([\d\.-]+)) (?P<lng>lon:([\d\.-]+)) (?P<heading>dir:([\d\.-]+)) (?P<speed>spd:([\d\.-]+)) (?P<lock>lck:([\d\.-]+)) (?P<time>time:([\d]+)) (?P<date>date:([\d]+)) (?P<status>trig:([\d]+))`)
  dataNames = dataRe.SubexpNames()
)

func (App *App) UpdateShuttles(dataFeed string, updateInterval int) {
  var st time.Duration
  for {

    // Sleep for n seconds before updating again
    time.Sleep(st)
    if st == 0 {
      // Initialize the sleep timer after the first sleep.  This lets us sleep during errors
      // when we 'continue' back to the top of the loop without waiting to sleep for the first
      // update run.
      st = time.Duration(updateInterval) * time.Second
    }

    // Make request to our tracking data feed
    resp, err := http.Get(dataFeed)
    if err != nil {
      log.Printf("error getting data feed: %v", err)
      continue;
    }
    defer resp.Body.Close()

    // Read response body content
    body, err := ioutil.ReadAll(resp.Body)
    if err != nil {
      log.Printf("error reading data feed: %v", err)
      continue;
    }

    delim := "eof"
    // Iterate through all vehicles returned by data feed
    vehicles_data := strings.Split(string(body), delim)
    // TODO: Figure out if this handles == 1 vehicle correctly or always assumes > 1.
    if len(vehicles_data) <= 1 {
      log.Printf("found no vehicles delineated by '%s'", delim)
    }

    updated := 0
    for i := 1; i < len(vehicles_data)-1; i++ {
      match := dataRe.FindAllStringSubmatch(vehicles_data[i], -1)[0]

      // Store named capturing group and matching expression as a key value pair
      result := map[string]string{}
      for i, item := range match {
        result[dataNames[i]] = item
      }

      // Create new vehicle update & insert update into database
      update := VehicleUpdate {
        Id:        bson.NewObjectId(),
        VehicleId: strings.Replace(result["id"], "Vehicle ID:", "", -1),
        Lat:       strings.Replace(result["lat"], "lat:", "", -1),
        Lng:       strings.Replace(result["lng"], "lon:", "", -1),
        Heading:   strings.Replace(result["heading"], "dir:", "", -1),
        Speed:     strings.Replace(result["speed"], "spd:", "", -1),
        Lock:      strings.Replace(result["lock"], "lck:", "", -1),
        Time:      strings.Replace(result["time"], "time:", "", -1),
        Date:      strings.Replace(result["date"], "date:", "", -1),
        Status:    strings.Replace(result["status"], "trig:", "", -1),
        Created:   time.Now()}

      if err := App.Updates.Insert(&update); err != nil {
        log.Printf("error inserting vehicle update(%v): %v", update, err)
      } else {
        updated++
      }
    }
    log.Printf("sucessfully updated %d/%d vehicles", updated, len(vehicles_data)-1)
  }
}

/**
 *  Route handlers - API requests,
 *                   serve view files
 */

type Vehicle struct {
  Id          bson.ObjectId                 `bson:"_id,omitempty"`
  VehicleId   string     `json:"vehicleId"   bson:"vehicleId,omitempty"`
  VehicleName string     `json:"vehicleName" bson:"vehicleName"`
  Created     time.Time  `json:"created"     bson:"created"`
}

/**
 * Find all vehicles in the database
 *
 */
func (App *App) VehiclesHandler(w http.ResponseWriter, r *http.Request) {
  // Find all vehicles in database
  var vehicles []Vehicle
  err := App.Vehicles.Find(bson.M{}).All(&vehicles)
  // Handle query errors
  if err != nil {
    http.Error(w, err.Error(), http.StatusInternalServerError)
  }
  // Send each vehicle to client as JSON
  vehiclesJSON, err := json.MarshalIndent(vehicles, "", " ")
  fmt.Fprintf(w, string(vehiclesJSON))
}

/**
 * Add a new vehicle to the database
 *
 */
func (App *App) VehiclesCreateHandler(w http.ResponseWriter, r *http.Request) {
  // Create new vehicle object using request fields
  vehicle := Vehicle{}
  vehicleData := json.NewDecoder(r.Body)
  err := vehicleData.Decode(&vehicle)
  // Error handling
  if err != nil {
    http.Error(w, err.Error(), http.StatusInternalServerError)
  }
  // Store new vehicle under vehicles collection
  err = App.Vehicles.Insert(&vehicle)
  // Error handling
  if err != nil {
    http.Error(w, err.Error(), http.StatusInternalServerError)
  }
}

/**
 * Vehicle Updates - Get most recent update for each
 *                   vehicle in the vehicles collection
 */
func (App *App) UpdatesHandler(w http.ResponseWriter, r *http.Request) {
  // Store updates for each vehicle
  var vehicles []Vehicle
  var updates []VehicleUpdate
  var update VehicleUpdate
  // Query all Vehicles
  err := App.Vehicles.Find(bson.M{}).All(&vehicles)
  // Handle errors
  if err != nil {
    http.Error(w, err.Error(), http.StatusInternalServerError)
  }
  // Find recent updates for each vehicle
  for _,vehicle := range vehicles {
    err := App.Updates.Find(bson.M{"vehicleId": vehicle.VehicleId}).Sort("-created").Limit(1).One(&update)

    if err == nil {
      updates = append(updates, update)
    }
  }
  // Convert updates to JSON
  u, err := json.MarshalIndent(updates, "", " ")
  // Handle JSON parsing errors
  if err != nil {
    http.Error(w, err.Error(), http.StatusInternalServerError)
  }
  // Send updates to client
  fmt.Fprint(w, string(u))
}

func IndexHandler(w http.ResponseWriter, r *http.Request) {
  http.ServeFile(w, r, "index.html")
}

func AdminHandler(w http.ResponseWriter, r *http.Request) {
  http.ServeFile(w, r, "admin.html")
}

/**
 *  Main - connect to database,
 *         handle routing,
 *         start tracker go routine,
 *         serve requests
 */

type Configuration struct {
  DataFeed        string
  UpdateInterval  int
  MongoUrl        string
  MongoPort       string
}

type App struct {
  Session    *mgo.Session
  Updates    *mgo.Collection
  Vehicles   *mgo.Collection
}

func ReadConfiguration(fileName string) (*Configuration, error) {
  // Open config file and decode JSON to Configuration struct
  file, err := os.Open(fileName)
  if err != nil {
    return nil, err
  }
  decoder := json.NewDecoder(file)
  config := Configuration{}
  if err := decoder.Decode(&config); err != nil {
    return nil, err
  }
  return &config, nil
}

func main() {
  // Read app configuration file
  config, err := ReadConfiguration("conf.json")
  if err != nil {
    log.Fatalf("error reading configuration file: %v", err)
  }

  // Connect to MongoDB
  session, err := mgo.Dial(config.MongoUrl + ":" + config.MongoPort)
  if err != nil {
    log.Fatalf("mongoDB connection failed: %v", err)
  }
  // close Mongo session when server terminates
  defer session.Close()

  // Create Shuttles object to store database session and collections
  App := &App{
    session,
    session.DB("shuttle_tracking").C("updates"),
    session.DB("shuttle_tracking").C("vehicles"),
  }

  // Start auto updater
  go App.UpdateShuttles(config.DataFeed, config.UpdateInterval)

  // Routing
  r := mux.NewRouter()
  r.HandleFunc("/", IndexHandler).Methods("GET")
  r.HandleFunc("/admin", AdminHandler).Methods("GET")
  r.HandleFunc("/admin/vehicles", AdminHandler).Methods("GET")
  r.HandleFunc("/admin/tracking", AdminHandler).Methods("GET")
  r.HandleFunc("/admin/stops", AdminHandler).Methods("GET")
  r.HandleFunc("/vehicles", App.VehiclesHandler).Methods("GET")
  r.HandleFunc("/vehicles/create", App.VehiclesCreateHandler).Methods("POST")
  r.HandleFunc("/updates", App.UpdatesHandler).Methods("GET")
  // Static files
  r.PathPrefix("/bower_components/").Handler(http.StripPrefix("/bower_components/", http.FileServer(http.Dir("bower_components/"))))
  r.PathPrefix("/static/").Handler(http.StripPrefix("/static/", http.FileServer(http.Dir("static/"))))
  // Serve requests
  http.Handle("/", r)
  if err := http.ListenAndServe(":8080", r); err != nil {
    log.Fatalf("Unable to ListenAndServe: %v", err)
  }
}
<|endoftext|>


<commit_before>// Copyright 2014 The Gogs Authors. All rights reserved.
// Copyright 2016 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// Gitea (git with a cup of tea) is a painless self-hosted Git Service.
package main // import "code.gitea.io/gitea"

import (
	"os"
	"strings"

	"code.gitea.io/gitea/cmd"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"

	// register supported doc types
	_ "code.gitea.io/gitea/modules/markup/csv"
	_ "code.gitea.io/gitea/modules/markup/markdown"
	_ "code.gitea.io/gitea/modules/markup/orgmode"

	"github.com/urfave/cli"
)

// Version holds the current Gitea version
var Version = "1.5.0-dev"

// Tags holds the build tags used
var Tags = ""

func init() {
	setting.AppVer = Version
	setting.AppBuiltWith = formatBuiltWith(Tags)
}

func main() {
	app := cli.NewApp()
	app.Name = "Gitea"
	app.Usage = "A painless self-hosted Git service"
	app.Description = `By default, gitea will start serving using the webserver with no
arguments - which can alternatively be run by running the subcommand web.`
	app.Version = Version + formatBuiltWith(Tags)
	app.Commands = []cli.Command{
		cmd.CmdWeb,
		cmd.CmdServ,
		cmd.CmdHook,
		cmd.CmdDump,
		cmd.CmdCert,
		cmd.CmdAdmin,
		cmd.CmdGenerate,
		cmd.CmdMigrate,
		cmd.CmdKeys,
	}
	app.Flags = append(app.Flags, cmd.CmdWeb.Flags...)
	app.Action = cmd.CmdWeb.Action
	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(4, "Failed to run app with %s: %v", os.Args, err)
	}
}

func formatBuiltWith(Tags string) string {
	if len(Tags) == 0 {
		return ""
	}

	return " built with: " + strings.Replace(Tags, " ", ", ", -1)
}
<commit_msg>Include Go toolchain to  --version (#5830)<commit_after>// Copyright 2014 The Gogs Authors. All rights reserved.
// Copyright 2016 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

// Gitea (git with a cup of tea) is a painless self-hosted Git Service.
package main // import "code.gitea.io/gitea"

import (
	"os"
	"runtime"
	"strings"

	"code.gitea.io/gitea/cmd"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"

	// register supported doc types
	_ "code.gitea.io/gitea/modules/markup/csv"
	_ "code.gitea.io/gitea/modules/markup/markdown"
	_ "code.gitea.io/gitea/modules/markup/orgmode"

	"github.com/urfave/cli"
)

// Version holds the current Gitea version
var Version = "1.5.0-dev"

// Tags holds the build tags used
var Tags = ""

func init() {
	setting.AppVer = Version
	setting.AppBuiltWith = formatBuiltWith(Tags)
}

func main() {
	app := cli.NewApp()
	app.Name = "Gitea"
	app.Usage = "A painless self-hosted Git service"
	app.Description = `By default, gitea will start serving using the webserver with no
arguments - which can alternatively be run by running the subcommand web.`
	app.Version = Version + formatBuiltWith(Tags)
	app.Commands = []cli.Command{
		cmd.CmdWeb,
		cmd.CmdServ,
		cmd.CmdHook,
		cmd.CmdDump,
		cmd.CmdCert,
		cmd.CmdAdmin,
		cmd.CmdGenerate,
		cmd.CmdMigrate,
		cmd.CmdKeys,
	}
	app.Flags = append(app.Flags, cmd.CmdWeb.Flags...)
	app.Action = cmd.CmdWeb.Action
	err := app.Run(os.Args)
	if err != nil {
		log.Fatal(4, "Failed to run app with %s: %v", os.Args, err)
	}
}

func formatBuiltWith(Tags string) string {
	if len(Tags) == 0 {
		return " built with " + runtime.Version()
	}

	return " built with " + runtime.Version() + " : " + strings.Replace(Tags, " ", ", ", -1)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

func main() {
	var (
		listen   = flag.String("listen", ":8080", "HTTP listen")
		interval = flag.Duration("interval", 1*time.Second, "broadcast interval")
		timeout  = flag.Duration("timeout", 5*time.Second, "websocket connection timeout")
	)
	flag.Parse()
	http.HandleFunc("/", handleSocket(newNumberStation(*interval), *timeout))
	log.Printf("listening on %s", *listen)
	log.Fatal(http.ListenAndServe(*listen, nil))
}

func handleSocket(s station, timeout time.Duration) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != "GET" {
			http.Error(w, "Method not allowed", 405)
			return
		}
		conn, err := websocket.Upgrade(w, r, nil, 1024, 1024)
		if _, ok := err.(websocket.HandshakeError); ok {
			http.Error(w, "Not a websocket handshake", 400)
			return
		} else if err != nil {
			log.Println(err)
			return
		}
		s.Subscribe(conn)
		wait(conn, timeout)
	}
}

type station interface {
	Subscribe(*websocket.Conn)
}

type numberStation struct {
	sync.RWMutex
	sub   chan *websocket.Conn
	unsub chan *websocket.Conn
}

func newNumberStation(interval time.Duration) *numberStation {
	s := &numberStation{
		sub:   make(chan *websocket.Conn),
		unsub: make(chan *websocket.Conn),
	}
	go s.loop(interval)
	return s
}

func (s *numberStation) loop(interval time.Duration) {
	m := map[*websocket.Conn]struct{}{}
	tick := time.Tick(interval)
	for {
		select {
		case <-tick:
			if len(m) <= 0 {
				continue
			}
			log.Printf("broadcasting to %d", len(m))
			number := rand.Intn(255)
			for conn := range m {
				go send(conn, number, s.unsub)
			}
		case conn := <-s.sub:
			log.Printf("%s: subscribed", conn.RemoteAddr())
			m[conn] = struct{}{}
		case conn := <-s.unsub:
			log.Printf("%s: unsubscribed", conn.RemoteAddr())
			delete(m, conn)
		}
	}
}

func (s *numberStation) Subscribe(conn *websocket.Conn) {
	s.sub <- conn
}

func send(conn *websocket.Conn, number int, unsub chan<- *websocket.Conn) {
	if err := conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf("%02x", number))); err != nil {
		log.Printf("%s: %s", conn.RemoteAddr(), err)
		unsub <- conn
	}
}

func wait(conn *websocket.Conn, timeout time.Duration) {
	defer conn.Close()
	for {
		if _, _, err := conn.ReadMessage(); err != nil {
			break
		}
	}
}
<commit_msg>Remove unused 'timeout'<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

func main() {
	var (
		listen   = flag.String("listen", ":8080", "HTTP listen")
		interval = flag.Duration("interval", 1*time.Second, "broadcast interval")
	)
	flag.Parse()
	http.HandleFunc("/", handleSocket(newNumberStation(*interval)))
	log.Printf("listening on %s", *listen)
	log.Fatal(http.ListenAndServe(*listen, nil))
}

func handleSocket(s station) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != "GET" {
			http.Error(w, "Method not allowed", 405)
			return
		}
		conn, err := websocket.Upgrade(w, r, nil, 1024, 1024)
		if _, ok := err.(websocket.HandshakeError); ok {
			http.Error(w, "Not a websocket handshake", 400)
			return
		} else if err != nil {
			log.Println(err)
			return
		}
		s.Subscribe(conn)
		wait(conn)
	}
}

type station interface {
	Subscribe(*websocket.Conn)
}

type numberStation struct {
	sync.RWMutex
	sub   chan *websocket.Conn
	unsub chan *websocket.Conn
}

func newNumberStation(interval time.Duration) *numberStation {
	s := &numberStation{
		sub:   make(chan *websocket.Conn),
		unsub: make(chan *websocket.Conn),
	}
	go s.loop(interval)
	return s
}

func (s *numberStation) loop(interval time.Duration) {
	m := map[*websocket.Conn]struct{}{}
	tick := time.Tick(interval)
	for {
		select {
		case <-tick:
			if len(m) <= 0 {
				continue
			}
			log.Printf("broadcasting to %d", len(m))
			number := rand.Intn(255)
			for conn := range m {
				go send(conn, number, s.unsub)
			}
		case conn := <-s.sub:
			log.Printf("%s: subscribed", conn.RemoteAddr())
			m[conn] = struct{}{}
		case conn := <-s.unsub:
			log.Printf("%s: unsubscribed", conn.RemoteAddr())
			delete(m, conn)
		}
	}
}

func (s *numberStation) Subscribe(conn *websocket.Conn) {
	s.sub <- conn
}

func send(conn *websocket.Conn, number int, unsub chan<- *websocket.Conn) {
	if err := conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf("%02x", number))); err != nil {
		log.Printf("%s: %s", conn.RemoteAddr(), err)
		unsub <- conn
	}
}

func wait(conn *websocket.Conn) {
	defer conn.Close()
	for {
		if _, _, err := conn.ReadMessage(); err != nil {
			break
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"

	_ "github.com/go-sql-driver/mysql"
	"github.com/tedsuo/ifrit"

	"github.com/cloudfoundry-incubator/cf-lager"
	"github.com/pivotal-cf-experimental/cf-mysql-quota-enforcer/config"
	"github.com/pivotal-cf-experimental/cf-mysql-quota-enforcer/database"
	"github.com/pivotal-cf-experimental/cf-mysql-quota-enforcer/enforcer"
	"github.com/pivotal-golang/lager"
)

type Config struct {
	Host         string
	Port         int
	User         string
	Password     string
	BrokerDBName string
}

func main() {
	flags := flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	runOnce := flags.Bool("runOnce", false, "Run only once instead of continuously")
	configFile := flags.String("configFile", "", "Location of config file")
	cf_lager.AddFlags(flags)
	flags.Parse(os.Args[1:])
	logger, _ := cf_lager.New("Quota Enforcer")

	logger.Info("Config file", lager.Data{"configFile": configFile})
	config, err := config.Load(*configFile)
	if err != nil {
		logger.Fatal("Failed to load config file", err)
	}

	brokerDBName := config.DBName
	if brokerDBName == "" {
		logger.Fatal("Must specify DBName in the config file", nil)
	}

	db, err := database.NewConnection(*config)
	if db != nil {
		defer db.Close()
	}

	if err != nil {
		logger.Fatal("Failed to open database connection", err)
	}

	logger.Info(
		"Database connection established.",
		lager.Data{
			"Host":         config.Host,
			"Port":         config.Port,
			"User":         config.User,
			"DatabaseName": brokerDBName,
		})

	violatorRepo := database.NewViolatorRepo(brokerDBName, db, logger)
	reformerRepo := database.NewReformerRepo(brokerDBName, db, logger)

	e := enforcer.NewEnforcer(violatorRepo, reformerRepo, logger)
	r := enforcer.NewRunner(e, logger)

	if *runOnce {
		logger.Info("Running once")

		err := e.EnforceOnce()
		if err != nil {
			logger.Info(fmt.Sprintf("Quota Enforcing Failed: %s", err.Error()))
		}
	} else {
		process := ifrit.Invoke(r)
		logger.Info("Running continuously")

		err := <-process.Wait()
		if err != nil {
			logger.Fatal("Quota Enforcing Failed", err)
		}
	}
}

func enforce(e enforcer.Enforcer, logger lager.Logger) {
	logger.Info("Enforcing")
	err := e.EnforceOnce()
	if err != nil {
		logger.Info(fmt.Sprintf("Enforcing Failed: %s", err.Error()))
	}
}
<commit_msg>Add default location of config file to be ./config.yml<commit_after>package main

import (
	"flag"
	"fmt"
	"os"

	_ "github.com/go-sql-driver/mysql"
	"github.com/tedsuo/ifrit"

	"github.com/cloudfoundry-incubator/cf-lager"
	"github.com/pivotal-cf-experimental/cf-mysql-quota-enforcer/config"
	"github.com/pivotal-cf-experimental/cf-mysql-quota-enforcer/database"
	"github.com/pivotal-cf-experimental/cf-mysql-quota-enforcer/enforcer"
	"github.com/pivotal-golang/lager"
)

type Config struct {
	Host         string
	Port         int
	User         string
	Password     string
	BrokerDBName string
}

func main() {
	flags := flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	runOnce := flags.Bool("runOnce", false, "Run only once instead of continuously")
	configFile := flags.String("configFile", "config.yml", "Location of config file")
	cf_lager.AddFlags(flags)
	flags.Parse(os.Args[1:])
	logger, _ := cf_lager.New("Quota Enforcer")

	logger.Info("Config file", lager.Data{"configFile": configFile})
	config, err := config.Load(*configFile)
	if err != nil {
		logger.Fatal("Failed to load config file", err)
	}

	brokerDBName := config.DBName
	if brokerDBName == "" {
		logger.Fatal("Must specify DBName in the config file", nil)
	}

	db, err := database.NewConnection(*config)
	if db != nil {
		defer db.Close()
	}

	if err != nil {
		logger.Fatal("Failed to open database connection", err)
	}

	logger.Info(
		"Database connection established.",
		lager.Data{
			"Host":         config.Host,
			"Port":         config.Port,
			"User":         config.User,
			"DatabaseName": brokerDBName,
		})

	violatorRepo := database.NewViolatorRepo(brokerDBName, db, logger)
	reformerRepo := database.NewReformerRepo(brokerDBName, db, logger)

	e := enforcer.NewEnforcer(violatorRepo, reformerRepo, logger)
	r := enforcer.NewRunner(e, logger)

	if *runOnce {
		logger.Info("Running once")

		err := e.EnforceOnce()
		if err != nil {
			logger.Info(fmt.Sprintf("Quota Enforcing Failed: %s", err.Error()))
		}
	} else {
		process := ifrit.Invoke(r)
		logger.Info("Running continuously")

		err := <-process.Wait()
		if err != nil {
			logger.Fatal("Quota Enforcing Failed", err)
		}
	}
}

func enforce(e enforcer.Enforcer, logger lager.Logger) {
	logger.Info("Enforcing")
	err := e.EnforceOnce()
	if err != nil {
		logger.Info(fmt.Sprintf("Enforcing Failed: %s", err.Error()))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"image"
	"image/draw"
	"image/jpeg"
	"image/png"
	"io"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"
)

type job struct {
	rowNumber int
	record    []string
}

var wg sync.WaitGroup

// main func
func main() {
	now := time.Now()

	file, err := os.Open("config.json")
	if err != nil {
		panic(err)
	}

	decoder := json.NewDecoder(file)
	config := Config{}

	err = decoder.Decode(&config)
	if err != nil {
		panic(err)
	}

	file, err = os.Open(config.CsvPath)
	if err != nil {
		panic(err)
	}

	reader := csv.NewReader(bufio.NewReader(file))
	record, err := reader.Read()

	rowNumber := 2

	jobs := make(chan job, 100)

	for w := 1; w <= 25; w++ {
		go worker(jobs)
	}

	for {
		record, err = reader.Read()

		// only for debug
		if err == io.EOF {
			break
		}

		wg.Add(1)
		// go mergeImages(rowNumber, record[26], record[5])
		jobs <- job{rowNumber, record}
		rowNumber++
	}

	wg.Wait()

	fmt.Printf("%s\n", time.Since(now))
}

// worker func
func worker(jobs <-chan job) {
	for j := range jobs {
		mergeImages(j.rowNumber, j.record[26], j.record[5])
	}
}

// mergeImages func
func mergeImages(rowNumber int, baseURL string, logoURL string) {
	getImages(baseURL, logoURL, rowNumber)
	glueImages(rowNumber)
}

// getImage func
func getImages(baseRUL string, logoURL string, rowNumber int) {
	getBaseImage(baseRUL, rowNumber)
	getLogoImage(logoURL, rowNumber)
}

// glueImages func
func glueImages(rowNumber int) {
	var basePath bytes.Buffer
	var logoPath bytes.Buffer
	var mergedPath bytes.Buffer

	basePath.WriteString("tmp/base/")
	basePath.WriteString(strconv.Itoa(rowNumber))

	bf, err := os.Open(basePath.String())
	if err != nil {
		panic(err)
	}

	baseImg, err := png.Decode(bufio.NewReader(bf))
	if err != nil {
		panic(err)
	}

	base := image.NewRGBA(image.Rect(baseImg.Bounds().Min.X, baseImg.Bounds().Min.Y, baseImg.Bounds().Max.X, baseImg.Bounds().Max.Y))

	logoPath.WriteString("tmp/logo/")
	logoPath.WriteString(strconv.Itoa(rowNumber))

	lf, err := os.Open(logoPath.String())
	if err != nil {
		panic(err)
	}

	logoImg, err := jpeg.Decode(bufio.NewReader(lf))
	if err != nil {
		panic(err)
	}

	draw.Draw(base, baseImg.Bounds(), baseImg, image.Point{0, 0}, draw.Src)
	draw.Draw(base, logoImg.Bounds(), logoImg, image.Point{15, 15}, draw.Src)

	mergedPath.WriteString("images/")
	mergedPath.WriteString(strconv.Itoa(rowNumber))
	mergedPath.WriteString(".png")

	mergeFile, err := os.Create(mergedPath.String())
	if err != nil {
		panic(err)
	}

	w := bufio.NewWriter(mergeFile)
	png.Encode(w, base)

	defer os.Remove(basePath.String())
	defer os.Remove(logoPath.String())

	fmt.Printf("%s %s\n", "Generated", mergedPath.String())

	wg.Done()
}

// getBaseImage func
func getBaseImage(baseRUL string, rowNumber int) {
	var imagePath bytes.Buffer
	imagePath.WriteString("tmp/base/")
	imagePath.WriteString(strconv.Itoa(rowNumber))

	out, err := os.Create(imagePath.String())
	if err != nil {
		panic(err)
	}

	response, err := http.Get(baseRUL)
	if err != nil {
		panic(err)
	}

	_, err = io.Copy(out, response.Body)
	if err != nil {
		panic(err)
	}
}

// getLogoImage func
func getLogoImage(logoURL string, rowNumber int) {
	var imagePath bytes.Buffer
	imagePath.WriteString("tmp/logo/")
	imagePath.WriteString(strconv.Itoa(rowNumber))

	out, err := os.Create(imagePath.String())
	if err != nil {
		panic(err)
	}

	response, err := http.Get(logoURL)
	if err != nil {
		panic(err)
	}

	_, err = io.Copy(out, response.Body)
	if err != nil {
		panic(err)
	}
}
<commit_msg>Unifying functions<commit_after>package main

import (
	"bufio"
	"bytes"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"image"
	"image/draw"
	"image/jpeg"
	"image/png"
	"io"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"
)

type job struct {
	rowNumber int
	record    []string
}

var wg sync.WaitGroup

// main func
func main() {
	now := time.Now()

	file, err := os.Open("config.json")
	if err != nil {
		panic(err)
	}

	decoder := json.NewDecoder(file)
	config := Config{}

	err = decoder.Decode(&config)
	if err != nil {
		panic(err)
	}

	file, err = os.Open(config.CsvPath)
	if err != nil {
		panic(err)
	}

	reader := csv.NewReader(bufio.NewReader(file))
	record, err := reader.Read()

	rowNumber := 2

	jobs := make(chan job, 100)

	for w := 1; w <= 25; w++ {
		go worker(jobs)
	}

	for {
		record, err = reader.Read()

		// only for debug
		if err == io.EOF {
			break
		}

		wg.Add(1)
		// go mergeImages(rowNumber, record[26], record[5])
		jobs <- job{rowNumber, record}
		rowNumber++
	}

	wg.Wait()

	fmt.Printf("%s\n", time.Since(now))
}

// worker func
func worker(jobs <-chan job) {
	for j := range jobs {
		mergeImages(j.rowNumber, j.record[26], j.record[5])
	}
}

// mergeImages func
func mergeImages(rowNumber int, baseURL string, logoURL string) {
	getImages(baseURL, logoURL, rowNumber)
	glueImages(rowNumber)
}

// getImage func
func getImages(baseURL string, logoURL string, rowNumber int) {
	getImage(baseURL, rowNumber, "base")
	getImage(logoURL, rowNumber, "logo")
}

// glueImages func
func glueImages(rowNumber int) {
	var basePath bytes.Buffer
	var logoPath bytes.Buffer
	var mergedPath bytes.Buffer

	basePath.WriteString("tmp/base/")
	basePath.WriteString(strconv.Itoa(rowNumber))

	bf, err := os.Open(basePath.String())
	if err != nil {
		panic(err)
	}

	baseImg, err := png.Decode(bufio.NewReader(bf))
	if err != nil {
		panic(err)
	}

	base := image.NewRGBA(image.Rect(baseImg.Bounds().Min.X, baseImg.Bounds().Min.Y, baseImg.Bounds().Max.X, baseImg.Bounds().Max.Y))

	logoPath.WriteString("tmp/logo/")
	logoPath.WriteString(strconv.Itoa(rowNumber))

	lf, err := os.Open(logoPath.String())
	if err != nil {
		panic(err)
	}

	logoImg, err := jpeg.Decode(bufio.NewReader(lf))
	if err != nil {
		panic(err)
	}

	draw.Draw(base, baseImg.Bounds(), baseImg, image.Point{0, 0}, draw.Src)
	draw.Draw(base, logoImg.Bounds(), logoImg, image.Point{15, 15}, draw.Src)

	mergedPath.WriteString("images/")
	mergedPath.WriteString(strconv.Itoa(rowNumber))
	mergedPath.WriteString(".png")

	mergeFile, err := os.Create(mergedPath.String())
	if err != nil {
		panic(err)
	}

	w := bufio.NewWriter(mergeFile)
	png.Encode(w, base)

	defer os.Remove(basePath.String())
	defer os.Remove(logoPath.String())

	fmt.Printf("%s %s\n", "Generated", mergedPath.String())

	wg.Done()
}

// getBaseImage func
func getImage(URL string, rowNumber int, imageType string) {
	var imagePath bytes.Buffer
	imagePath.WriteString("tmp/")
	imagePath.WriteString(imageType)
	imagePath.WriteString("/")
	imagePath.WriteString(strconv.Itoa(rowNumber))

	out, err := os.Create(imagePath.String())
	if err != nil {
		panic(err)
	}

	response, err := http.Get(URL)
	if err != nil {
		panic(err)
	}

	_, err = io.Copy(out, response.Body)
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/jrzimmerman/bestrida-server-go/handlers"
	"github.com/jrzimmerman/bestrida-server-go/models"
	"github.com/jrzimmerman/bestrida-server-go/utils"
)

func main() {
	port := utils.GetEnvString("PORT")
	mux := handlers.API()
	// close DB connection
	defer models.Close()

	srv := &http.Server{
		Addr:    ":" + port,
		Handler: mux,
	}

	// We want to report the listener is closed.
	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		log.Printf("Listening on http://localhost:%s\n", port)

		srv.ListenAndServe()
		wg.Done()
	}()

	// Listen for an interrupt signal from the OS.
	osSignals := make(chan os.Signal)
	signal.Notify(osSignals, os.Interrupt)

	// Wait for a signal to shutdown.
	<-osSignals

	log.Println("\nShutting down the server...")

	// Create a context to attempt a graceful 5 second shutdown.
	const timeout = 5 * time.Second
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	// Attempt the graceful shutdown by closing the listener and
	// completing all inflight requests.
	if err := srv.Shutdown(ctx); err != nil {
		log.Printf("Graceful shutdown did not complete in %v : %v", timeout, err)

		// Looks like we timedout on the graceful shutdown. Kill it hard.
		if err := srv.Close(); err != nil {
			log.Printf("Error killing server : %v", err)
		}
	}

	// Wait for the listener to report it is closed.
	wg.Wait()
	log.Println("Server gracefully stopped")
}
<commit_msg>use buffered channel to wait for os signal<commit_after>package main

import (
	"context"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/jrzimmerman/bestrida-server-go/handlers"
	"github.com/jrzimmerman/bestrida-server-go/models"
	"github.com/jrzimmerman/bestrida-server-go/utils"
)

func main() {
	port := utils.GetEnvString("PORT")
	mux := handlers.API()
	// close DB connection
	defer models.Close()

	srv := &http.Server{
		Addr:    ":" + port,
		Handler: mux,
	}

	// We want to report the listener is closed.
	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		log.Printf("Listening on http://localhost:%s\n", port)

		srv.ListenAndServe()
		wg.Done()
	}()

	// Listen for an interrupt signal from the OS. Use a buffered
	// channel because of how the signal package is implemented.
	osSignals := make(chan os.Signal, 1)
	signal.Notify(osSignals, os.Interrupt)

	// Wait for a signal to shutdown.
	<-osSignals

	// Create a context to attempt a graceful 5 second shutdown.
	const timeout = 5 * time.Second
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	log.Println("\nShutting down the server...")

	// Attempt the graceful shutdown by closing the listener and
	// completing all inflight requests.
	if err := srv.Shutdown(ctx); err != nil {
		log.Printf("shutdown : Graceful shutdown did not complete in %v : %v", timeout, err)

		// Looks like we timedout on the graceful shutdown. Kill it hard.
		if err := srv.Close(); err != nil {
			log.Printf("shutdown : Error killing server : %v", err)
		}
	}

	// Wait for the listener to report it is closed.
	wg.Wait()
	log.Println("Server gracefully stopped")
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"compress/gzip"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/uttamgandhi24/whisper-go/whisper"
	"io/ioutil"
	"log"
	"math"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
)

type MigrationData struct {
	wspFile         string
	relativePath    string
	exportFileName  string
	measurement     string
	tags            string
	field           string
	matched         bool
	exported        bool
}

type TagKeyValue struct {
	Tagkey   string `json:"tagkey"`
	Tagvalue string `json:"tagvalue"`
}

type MigrationConfig struct {
	Pattern     string        `json:"pattern"`
	Measurement string        `json:"measurement"`
	Tags        []TagKeyValue `json:"tags"`
	Field       string        `json:"field"`
}

// Global vars
var migrationConfig []MigrationConfig
var exportedFileNumber = 0
var (
	verbose           = flag.Bool("verbose", false, "Configuration file for measurement and tags.")
	wspPath           = flag.String("wsp-path", "", "Whisper files folder path.")
	exportPath        = flag.String("export-path", "", "Directory to export line protocol files.")
	configFile        = flag.String("config-file", "", "Configuration file for measurement and tags.")
	fromFlag          = flag.Uint("from", 0, "Configuration file for measurement and tags.")
	untilFlag         = flag.Uint("until", uint(^uint32(0)), "Configuration file for measurement and tags.")
	gzipped           = flag.Bool("gz", false, "Export data in a gzipped file")
)

func main() {
	flag.Parse()

	// List files
	fileList := []string{}
	listWspFiles(&fileList, *wspPath)

	// Open migration config file
	loadConfigFile(&migrationConfig, *configFile)

	// Go through wsp files and figure out tags, measurements, file names, etc.
	var migrationData []MigrationData
	for _, wspFile := range fileList {
		data := MigrationData{}
		data.wspFile = wspFile
		data.relativePath = strings.TrimPrefix(wspFile, *wspPath)

		// Figure exported filename
		exportedFileNumber += 1
		data.exportFileName = fmt.Sprintf("%08d.txt", exportedFileNumber)

		// Add the .gz if the export needs to be compressed
		if *gzipped {
			data.exportFileName += ".gz"
		}

		// Assign the right measurment, field and tags
		data.assignConfig()

		if data.matched {
			migrationData = append(migrationData, data)
		} else {
			fmt.Println("File didn't match any config patterns: ", data.wspFile)
		}
	}

	// Time boundaries
	var from uint32 = uint32(*fromFlag)
	var until uint32 = uint32(*untilFlag)

	// Warning starting exporting
	fmt.Println("----------------")
	fmt.Println("Exporting", len(migrationData), "series to", *exportPath)
	if askForConfirmation("Proceed ?") == false {
		os.Exit(1)
	}
	fmt.Println("----------------")

	// Go through wsp files and export data
	for k, migration := range migrationData {
		// Open whisper file with driver
		w, err := whisper.Open(migration.wspFile)
		check(err)

		for i, archive := range w.Header.Archives {
			var wspPoints []whisper.Point

			// Give a name to the retention
			retentionName := fmt.Sprintf("%d-%d", archive.SecondsPerPoint, archive.Points)

			// Go through points
			points, err := w.DumpArchive(i)
			if err != nil {
				if *verbose {
					fmt.Println("Unexpected EOF:", retentionName, "from", migration.wspFile)
				}
				continue
			}
			for _, point := range points {
				// Skip the point on certain conditions
				if point.Value == 0 {
					continue
				}
				if point.Timestamp < from || point.Timestamp > until {
					continue
				}

				wspPoints = append(wspPoints, point)
			}

			// If there is nothing to write, skip file creation
			if len(wspPoints) == 0 {
				if *verbose {
					fmt.Println("Skipped:", retentionName, "from", migration.wspFile)
				}
				continue
			}

			// Makes sure the directory exists
			path := *exportPath + "/" + retentionName
			os.MkdirAll(path ,0755);

			// Write file
			func() {
				// Open file and prepare writer
				f, err := os.Create(path + "/" + migration.exportFileName)
				check(err)
				defer f.Close()

				var gf *gzip.Writer
				var fw *bufio.Writer
				if *gzipped {
					gf = gzip.NewWriter(f)
					fw = bufio.NewWriter(gf)
					defer gf.Close()
				} else {
					fw = bufio.NewWriter(f)
				}

				// Print all points
				for _, point := range wspPoints {
					line := migration.lineprotocol(point, archive.SecondsPerPoint) + "\n"
					_, err := fw.WriteString(line)
					check(err)
				}

				// Flush writer and close file
				fw.Flush()
			}() // <-- function call to defer

			if *verbose {
				fmt.Println("Exported:", path + "/" + migration.exportFileName, "from", migration.wspFile)
			} else {
				fmt.Printf("\rExported: %2d", k)
			}
		}

		w.Close()
	}
}

// Check errors
func check(e error) {
	if e != nil {
		panic(e)
	}
}

// Source: https://gist.github.com/m4ng0squ4sh/3dcbb0c8f6cfe9c66ab8008f55f8f28b
// askForConfirmation asks the user for confirmation. A user must type in "yes" or "no" and
// then press enter. It has fuzzy matching, so "y", "Y", "yes", "YES", and "Yes" all count as
// confirmations. If the input is not recognized, it will ask again. The function does not return
// until it gets a valid response from the user.
func askForConfirmation(s string) bool {
	reader := bufio.NewReader(os.Stdin)

	for {
		fmt.Printf("%s [y/n]: ", s)

		response, err := reader.ReadString('\n')
		if err != nil {
			log.Fatal(err)
		}

		response = strings.ToLower(strings.TrimSpace(response))

		if response == "y" || response == "yes" {
			return true
		} else if response == "n" || response == "no" {
			return false
		}
	}
}

// Create the list of wsp files
func listWspFiles(fileList *[]string, searchDir string) {
	err := filepath.Walk(searchDir, func(path string, f os.FileInfo, err error) error {
		if os.IsNotExist(err) { //search dir does not exist
			return nil
		}
		// Only add wsp files to the list
		if strings.HasSuffix(f.Name(), "wsp") {
			*fileList = append(*fileList, path)
		}
		return nil
	})
	if err != nil {
		fmt.Println("Error listing files:")
		fmt.Println(err)
	}
}

// Read the config file and populate migrartionData.tagConfigs
func loadConfigFile(migrationConfig *[]MigrationConfig, filename string) error {
	raw, err := ioutil.ReadFile(filename)
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}
	return json.Unmarshal(raw, &migrationConfig)
}

// Generate the influxdb line protocol string for a given point
func (migrationData *MigrationData) lineprotocol(point whisper.Point, factor uint32) string {
	var line string
	line += migrationData.measurement
	line += migrationData.tags
	line += " "
	line += migrationData.field + "=" + strconv.FormatFloat(math.Ceil(point.Value * float64(factor)), 'f', -1, 64)
	line += " "
	line += strconv.FormatInt(int64(point.Timestamp), 10)
	return line
}

// Get measurement, tags and field by matching the whisper filename with a
// pattern in the config file
// This part is inspired by the project https://github.com/influxdata/whisper-migrator
func (migrationData *MigrationData) assignConfig() {

	wspMeasurement := migrationData.wspFile
	wspMeasurement = strings.TrimSuffix(wspMeasurement, ".wsp")
	wspMeasurement = strings.Replace(wspMeasurement, "/", ".", -1)
	wspMeasurement = strings.Replace(wspMeasurement, ",", "_", -1)
	wspMeasurement = strings.Replace(wspMeasurement, " ", "_", -1)

	// Filename matching
	// TODO catch strings that don't match to the end
	var tagConfig MigrationConfig
	var matched []string
	var matchedArr [][]string
	var wildcards [][]string
	filenameMatched := false
	for _, tagConfig = range migrationConfig {
		reWild := regexp.MustCompile("{{\\s*(\\S+)\\s*}}")

		// Prepare regex pattern
		pattern := strings.Replace(tagConfig.Pattern, ".", "\\.", -1)
		pattern = reWild.ReplaceAllLiteralString(pattern, "([^.]+)")

		// List the matching values (Base and groups)
		re := regexp.MustCompile(pattern)

		matchedArr = re.FindAllStringSubmatch(wspMeasurement, -1)

		if matchedArr != nil {
			filenameMatched = true
			matched = matchedArr[0]
			// List of replacement wildcards like "{{ host }}"
			wildcards = reWild.FindAllStringSubmatch(tagConfig.Pattern, -1)
			break
		}
	}

	// Exit if there was no match
	if filenameMatched == false {
		return
	} else {
		migrationData.matched = true
	}

	// Fill the migrationData object
	migrationData.measurement = tagConfig.Measurement
	migrationData.field = tagConfig.Field
	for j := 0; j < len(tagConfig.Tags); j++ {
		migrationData.tags += ","
		migrationData.tags += tagConfig.Tags[j].Tagkey
		migrationData.tags += "="
		migrationData.tags += tagConfig.Tags[j].Tagvalue
	}

	// If measurement or field parameters is not defined
	if tagConfig.Measurement == "" {
		parts := strings.Split(wspMeasurement, ".")
		migrationData.measurement = parts[len(parts)-1]
	}
	if tagConfig.Field == "" {
		migrationData.field = "value"
	}

	// Replace "{{ wildcard }}" with matched values in order
	// (reversed to avoid overlapping of bigger numbers)
	// TODO issues if matched contains "$n"
	for i := len(matched) - 1; i > 0; i-- {
		re := regexp.MustCompile("{{\\s" + wildcards[i - 1][1] + "\\s}}")

		migrationData.measurement = re.ReplaceAllLiteralString(migrationData.measurement, matched[i])
		migrationData.field = re.ReplaceAllLiteralString(migrationData.field, matched[i])
		migrationData.tags = re.ReplaceAllLiteralString(migrationData.tags, matched[i])
	}
}
<commit_msg>Add -zeros flag<commit_after>package main

import (
	"bufio"
	"compress/gzip"
	"encoding/json"
	"flag"
	"fmt"
	"github.com/uttamgandhi24/whisper-go/whisper"
	"io/ioutil"
	"log"
	"math"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
)

type MigrationData struct {
	wspFile         string
	relativePath    string
	exportFileName  string
	measurement     string
	tags            string
	field           string
	matched         bool
	exported        bool
}

type TagKeyValue struct {
	Tagkey   string `json:"tagkey"`
	Tagvalue string `json:"tagvalue"`
}

type MigrationConfig struct {
	Pattern     string        `json:"pattern"`
	Measurement string        `json:"measurement"`
	Tags        []TagKeyValue `json:"tags"`
	Field       string        `json:"field"`
}

// Global vars
var migrationConfig []MigrationConfig
var exportedFileNumber = 0
var (
	verbose      = flag.Bool("verbose", false, "Configuration file for measurement and tags.")
	wspPath      = flag.String("wsp-path", "", "Whisper files folder path.")
	exportPath   = flag.String("export-path", "", "Directory to export line protocol files.")
	configFile   = flag.String("config-file", "", "Configuration file for measurement and tags.")
	fromFlag     = flag.Uint("from", 0, "Only export points after the given timestamp.")
	untilFlag    = flag.Uint("until", uint(^uint32(0)), "Only export points before the given timestamp.")
	gzipped      = flag.Bool("gz", false, "Export data in a gzipped file.")
	exportZeros    = flag.Bool("zeros", false, "Export null values (equal to zero). Those are ignored by default.")
)

func main() {
	flag.Parse()

	// List files
	fileList := []string{}
	listWspFiles(&fileList, *wspPath)

	// Open migration config file
	loadConfigFile(&migrationConfig, *configFile)

	// Go through wsp files and figure out tags, measurements, file names, etc.
	var migrationData []MigrationData
	for _, wspFile := range fileList {
		data := MigrationData{}
		data.wspFile = wspFile
		data.relativePath = strings.TrimPrefix(wspFile, *wspPath)

		// Figure exported filename
		exportedFileNumber += 1
		data.exportFileName = fmt.Sprintf("%08d.txt", exportedFileNumber)

		// Add the .gz if the export needs to be compressed
		if *gzipped {
			data.exportFileName += ".gz"
		}

		// Assign the right measurment, field and tags
		data.assignConfig()

		if data.matched {
			migrationData = append(migrationData, data)
		} else {
			fmt.Println("File didn't match any config patterns: ", data.wspFile)
		}
	}

	// Time boundaries
	var from uint32 = uint32(*fromFlag)
	var until uint32 = uint32(*untilFlag)

	// Warning starting exporting
	fmt.Println("----------------")
	fmt.Println("Exporting", len(migrationData), "series to", *exportPath)
	if askForConfirmation("Proceed ?") == false {
		os.Exit(1)
	}
	fmt.Println("----------------")

	// Go through wsp files and export data
	for k, migration := range migrationData {
		// Open whisper file with driver
		w, err := whisper.Open(migration.wspFile)
		check(err)

		for i, archive := range w.Header.Archives {
			var wspPoints []whisper.Point

			// Give a name to the retention
			retentionName := fmt.Sprintf("%d-%d", archive.SecondsPerPoint, archive.Points)

			// Go through points
			points, err := w.DumpArchive(i)
			if err != nil {
				if *verbose {
					fmt.Println("Unexpected EOF:", retentionName, "from", migration.wspFile)
				}
				continue
			}
			for _, point := range points {
				// Skip the point on certain conditions
				if !*exportZeros && point.Value == 0 {
					continue
				}
				if point.Timestamp < from || point.Timestamp > until {
					continue
				}

				wspPoints = append(wspPoints, point)
			}

			// If there is nothing to write, skip file creation
			if len(wspPoints) == 0 {
				if *verbose {
					fmt.Println("Skipped:", retentionName, "from", migration.wspFile)
				}
				continue
			}

			// Makes sure the directory exists
			path := *exportPath + "/" + retentionName
			os.MkdirAll(path ,0755);

			// Write file
			func() {
				// Open file and prepare writer
				f, err := os.Create(path + "/" + migration.exportFileName)
				check(err)
				defer f.Close()

				var gf *gzip.Writer
				var fw *bufio.Writer
				if *gzipped {
					gf = gzip.NewWriter(f)
					fw = bufio.NewWriter(gf)
					defer gf.Close()
				} else {
					fw = bufio.NewWriter(f)
				}

				// Print all points
				for _, point := range wspPoints {
					line := migration.lineprotocol(point, archive.SecondsPerPoint) + "\n"
					_, err := fw.WriteString(line)
					check(err)
				}

				// Flush writer and close file
				fw.Flush()
			}() // <-- function call to defer

			if *verbose {
				fmt.Println("Exported:", path + "/" + migration.exportFileName, "from", migration.wspFile)
			} else {
				fmt.Printf("\rExported: %2d", k)
			}
		}

		w.Close()
	}
}

// Check errors
func check(e error) {
	if e != nil {
		panic(e)
	}
}

// Source: https://gist.github.com/m4ng0squ4sh/3dcbb0c8f6cfe9c66ab8008f55f8f28b
// askForConfirmation asks the user for confirmation. A user must type in "yes" or "no" and
// then press enter. It has fuzzy matching, so "y", "Y", "yes", "YES", and "Yes" all count as
// confirmations. If the input is not recognized, it will ask again. The function does not return
// until it gets a valid response from the user.
func askForConfirmation(s string) bool {
	reader := bufio.NewReader(os.Stdin)

	for {
		fmt.Printf("%s [y/n]: ", s)

		response, err := reader.ReadString('\n')
		if err != nil {
			log.Fatal(err)
		}

		response = strings.ToLower(strings.TrimSpace(response))

		if response == "y" || response == "yes" {
			return true
		} else if response == "n" || response == "no" {
			return false
		}
	}
}

// Create the list of wsp files
func listWspFiles(fileList *[]string, searchDir string) {
	err := filepath.Walk(searchDir, func(path string, f os.FileInfo, err error) error {
		if os.IsNotExist(err) { //search dir does not exist
			return nil
		}
		// Only add wsp files to the list
		if strings.HasSuffix(f.Name(), "wsp") {
			*fileList = append(*fileList, path)
		}
		return nil
	})
	if err != nil {
		fmt.Println("Error listing files:")
		fmt.Println(err)
	}
}

// Read the config file and populate migrartionData.tagConfigs
func loadConfigFile(migrationConfig *[]MigrationConfig, filename string) error {
	raw, err := ioutil.ReadFile(filename)
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}
	return json.Unmarshal(raw, &migrationConfig)
}

// Generate the influxdb line protocol string for a given point
func (migrationData *MigrationData) lineprotocol(point whisper.Point, factor uint32) string {
	var line string
	line += migrationData.measurement
	line += migrationData.tags
	line += " "
	line += migrationData.field + "=" + strconv.FormatFloat(math.Ceil(point.Value * float64(factor)), 'f', -1, 64)
	line += " "
	line += strconv.FormatInt(int64(point.Timestamp), 10)
	return line
}

// Get measurement, tags and field by matching the whisper filename with a
// pattern in the config file
// This part is inspired by the project https://github.com/influxdata/whisper-migrator
func (migrationData *MigrationData) assignConfig() {

	wspMeasurement := migrationData.wspFile
	wspMeasurement = strings.TrimSuffix(wspMeasurement, ".wsp")
	wspMeasurement = strings.Replace(wspMeasurement, "/", ".", -1)
	wspMeasurement = strings.Replace(wspMeasurement, ",", "_", -1)
	wspMeasurement = strings.Replace(wspMeasurement, " ", "_", -1)

	// Filename matching
	// TODO catch strings that don't match to the end
	var tagConfig MigrationConfig
	var matched []string
	var matchedArr [][]string
	var wildcards [][]string
	filenameMatched := false
	for _, tagConfig = range migrationConfig {
		reWild := regexp.MustCompile("{{\\s*(\\S+)\\s*}}")

		// Prepare regex pattern
		pattern := strings.Replace(tagConfig.Pattern, ".", "\\.", -1)
		pattern = reWild.ReplaceAllLiteralString(pattern, "([^.]+)")

		// List the matching values (Base and groups)
		re := regexp.MustCompile(pattern)

		matchedArr = re.FindAllStringSubmatch(wspMeasurement, -1)

		if matchedArr != nil {
			filenameMatched = true
			matched = matchedArr[0]
			// List of replacement wildcards like "{{ host }}"
			wildcards = reWild.FindAllStringSubmatch(tagConfig.Pattern, -1)
			break
		}
	}

	// Exit if there was no match
	if filenameMatched == false {
		return
	} else {
		migrationData.matched = true
	}

	// Fill the migrationData object
	migrationData.measurement = tagConfig.Measurement
	migrationData.field = tagConfig.Field
	for j := 0; j < len(tagConfig.Tags); j++ {
		migrationData.tags += ","
		migrationData.tags += tagConfig.Tags[j].Tagkey
		migrationData.tags += "="
		migrationData.tags += tagConfig.Tags[j].Tagvalue
	}

	// If measurement or field parameters is not defined
	if tagConfig.Measurement == "" {
		parts := strings.Split(wspMeasurement, ".")
		migrationData.measurement = parts[len(parts)-1]
	}
	if tagConfig.Field == "" {
		migrationData.field = "value"
	}

	// Replace "{{ wildcard }}" with matched values in order
	// (reversed to avoid overlapping of bigger numbers)
	// TODO issues if matched contains "$n"
	for i := len(matched) - 1; i > 0; i-- {
		re := regexp.MustCompile("{{\\s" + wildcards[i - 1][1] + "\\s}}")

		migrationData.measurement = re.ReplaceAllLiteralString(migrationData.measurement, matched[i])
		migrationData.field = re.ReplaceAllLiteralString(migrationData.field, matched[i])
		migrationData.tags = re.ReplaceAllLiteralString(migrationData.tags, matched[i])
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"golang.org/x/net/context"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"

	skicka "github.com/pellaeon/skicka"
	"github.com/pellaeon/skicka/gdrive"

	"github.com/pellaeon/goas/v3/logger"
)

///////////////////////////////////////////////////////////////////////////
// Global Variables

const BLKSIZE = 4096

var progName = filepath.Base(os.Args[0])

type debugging bool

var (
	gd *gdrive.GDrive

	// The key is only set if encryption is needed (i.e. if -encrypt is
	// provided for an upload, or if an encrypted file is encountered
	// during 'download' or 'cat').
	key []byte

	debug   debugging
	verbose debugging
	quiet   bool

	// Configuration read in from the skicka config file.
	config struct {
		Google struct {
			ClientId     string
			ClientSecret string
			// If set, is appended to all http requests via ?key=XXX.
			ApiKey string
		}
		Encryption struct {
			Salt             string
			Passphrase_hash  string
			Encrypted_key    string
			Encrypted_key_iv string
		}
		Upload struct {
			Ignored_Regexp         []string
			Bytes_per_second_limit int
		}
		Download struct {
			Bytes_per_second_limit int
		}
	}

	// Various statistics gathered along the way. These all should be
	// updated using atomic operations since we often have multiple threads
	// working concurrently for uploads and downloads.
	stats struct {
		DiskReadBytes     int64
		DiskWriteBytes    int64
		UploadBytes       int64
		DownloadBytes     int64
		LocalFilesUpdated int64
		DriveFilesUpdated int64
	}

	// Smaller files will be handled with multiple threads going at once;
	// doing so improves bandwidth utilization since round-trips to the
	// Drive APIs take a while.  (However, we don't want too have too many
	// workers; this would both lead to lots of 403 rate limit errors...)
	nWorkers int
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", progName)
	fmt.Fprintf(os.Stderr, "  %s MOUNTPOINT\n", progName)
	flag.PrintDefaults()
}

func userHomeDir() string {
	if runtime.GOOS == "windows" {
		home := os.Getenv("HOMEDRIVE") + os.Getenv("HOMEPATH")
		if home == "" {
			home = os.Getenv("USERPROFILE")
		}
		return home
	}
	return os.Getenv("HOME")
}

func main() {
	l := logger.NewStandardTimeLogger(os.Stdout, time.StampMicro)
	logger.SetLogger(l)
	logger.SetLevel(logger.LevelDebug)
	log.SetFlags(0)
	log.SetPrefix(progName + ": ")

	// Initialize skicka
	home := userHomeDir()
	tokenCacheFilename := flag.String("tokencache",
		filepath.Join(home, ".skicka.tokencache.json"),
		"OAuth2 token cache file")
	configFilename := flag.String("config",
		filepath.Join(home, ".skicka.config"),
		"Configuration file")
	metadataCacheFilename := flag.String("metadata-cache-file",
		filepath.Join(home, "/.skicka.metadata.cache"),
		"Filename for local cache of Google Drive file metadata")
	nw := flag.Int("num-threads", 4, "Number of threads to use for uploads/downloads")
	vb := flag.Bool("verbose", false, "Enable verbose output")
	dbg := flag.Bool("debug", false, "Enable debugging output")
	qt := flag.Bool("quiet", false, "Suppress non-error messages")
	dumpHTTP := flag.Bool("dump-http", false, "Dump http traffic")
	flakyHTTP := flag.Bool("flaky-http", false, "Add flakiness to http traffic")
	noBrowserAuth := flag.Bool("no-browser-auth", false,
		"Don't try launching browser for authorization")
	flag.Usage = usage
	flag.Parse()

	var mountpoint string
	if flag.NArg() == 1 {
		mountpoint = flag.Arg(0)
	} else {
		mountpoint = *flag.String("mountpoint", "", "Mountpoint")
	}
	if mountpoint == "" {
		usage()
		os.Exit(2)
	}

	nWorkers = *nw

	debug = debugging(*dbg)
	verbose = debugging(*vb || bool(debug))
	quiet = *qt

	skicka.ReadConfigFile(*configFilename)

	// Set up the basic http.Transport.
	transport := http.DefaultTransport
	if tr, ok := transport.(*http.Transport); ok {
		// Increase the default number of open connections per destination host
		// to be enough for the number of goroutines we run concurrently for
		// uploads/downloads; this gives some benefit especially for uploading
		// small files.
		tr.MaxIdleConnsPerHost = 4
	} else {
		skicka.PrintErrorAndExit(fmt.Errorf("DefaultTransport not an *http.Transport?"))
	}
	if *flakyHTTP {
		transport = skicka.NewFlakyTransport(transport)
	}
	if *dumpHTTP {
		//transport = skicka.LoggingTransport{transport: transport}
		// TODO
	}

	// And now upgrade to the OAuth Transport *http.Client.
	client, err := skicka.GetOAuthClient(*tokenCacheFilename, !*noBrowserAuth,
		transport)
	if err != nil {
		skicka.PrintErrorAndExit(fmt.Errorf("error with OAuth2 Authorization: %v ", err))
	}

	// Choose the appropriate callback function for the GDrive object to
	// use for debugging output.
	var dpf func(s string, args ...interface{})
	if debug {
		dpf = skicka.DebugPrint
	} else {
		dpf = skicka.DebugNoPrint
	}

	gd, err := gdrive.New(config.Upload.Bytes_per_second_limit,
		config.Download.Bytes_per_second_limit, dpf, client,
		*metadataCacheFilename, quiet)
	if err != nil {
		log.Fatal(fmt.Errorf("error creating Google Drive "+
			"client: %v", err))
		os.Exit(3)
	}

	c, err := fuse.Mount(mountpoint)
	if err != nil {
		log.Fatalf("Mount failed: %v", err)
	}
	defer c.Close()

	filesys := &FS{
		gd: gd,
	}
	if err := fs.Serve(c, filesys); err != nil {
		log.Fatalf("Serve failed: %v", err)
	}

	// check if the mount process has an error to report
	<-c.Ready
	if err := c.MountError; err != nil {
		log.Fatalf("mount process error: %v", err)
	}
}

type FS struct {
	gd *gdrive.GDrive
}

var _ fs.FS = (*FS)(nil)

func (f *FS) Root() (fs.Node, error) {
	gd_file, err := f.gd.GetFile("/")
	if err != nil {
		logger.Errorf("FS Root() : %v", err)
	}
	return Dir{
		sk_file: gd_file,
		fs:      f,
	}, err
}

var _ fs.FSStatfser = (*FS)(nil)

func (fs *FS) Statfs(ctx context.Context, req *fuse.StatfsRequest, resp *fuse.StatfsResponse) error {
	usage, err := fs.gd.GetDriveUsage()
	if err != nil {
		logger.Errorf("Statfs() error: %v", err)
	}
	resp.Blocks = uint64(usage.Capacity / BLKSIZE)
	resp.Bfree = uint64((usage.Capacity - usage.Used) / BLKSIZE)
	resp.Bavail = uint64((usage.Capacity - usage.Used) / BLKSIZE)
	resp.Files = 9999 //XXX
	resp.Ffree = 9999 //XXX
	resp.Bsize = BLKSIZE
	resp.Namelen = 32767 //http://www.aurelp.com/2014/09/10/what-is-the-maximum-name-length-for-a-file-on-google-drive/
	resp.Frsize = BLKSIZE
	req.Respond(resp)
	return err
}

var _ fs.Node = (*Dir)(nil)

type Dir struct {
	sk_file *gdrive.File
	fs      *FS
}

func (n Dir) Attr(ctx context.Context, attr *fuse.Attr) error {
	var sum uint64
	for _, c := range n.sk_file.Id {
		sum += uint64(c)
	}
	attr.Inode = sum
	attr.Size = 4096 // XXX
	attr.Blocks = 0
	attr.Atime = n.sk_file.ModTime
	attr.Mtime = n.sk_file.ModTime
	attr.Ctime = n.sk_file.ModTime
	attr.Crtime = n.sk_file.ModTime
	attr.Mode = os.ModeDir | 0755
	attr.Nlink = 0

	return nil
}

type DirHandle struct {
	sk_file *gdrive.File
}

var _ fs.HandleReadDirAller = (*DirHandle)(nil)

func (dh *DirHandle) ReadDirAll(ctx context.Context) ([]fuse.Dirent, error) {
	logger.Debugf("ReadDirAll %s", dh.sk_file.Path)
	files_in_folder, err := gd.GetFilesInFolder(dh.sk_file.Path)
	logger.Debugf("ReadDirAll %+v", files_in_folder)
	if err != nil {
		log.Fatalf("ReadDirAll failed: %v", err)
	}
	var res []fuse.Dirent
	for _, file := range files_in_folder {
		var de fuse.Dirent
		de.Name = file.Path[strings.LastIndex(file.Path, "/")+1:]
		if file.IsFolder() {
			de.Type = fuse.DT_Dir
		} else {
			de.Type = fuse.DT_File
		}
		sum := uint64(0)
		for _, c := range file.Id {
			sum += uint64(c)
		}
		de.Inode = sum
		res = append(res, de)
	}
	return res, err
}

var _ fs.Node = (*File)(nil)

type File struct {
	sk_file *gdrive.File
	fs      *FS
}

func (n File) Attr(ctx context.Context, attr *fuse.Attr) error {
	sum := uint64(0)
	for _, c := range n.sk_file.Id {
		sum += uint64(c)
	}
	attr.Inode = sum
	attr.Size = uint64(n.sk_file.FileSize)
	attr.Blocks = uint64(n.sk_file.FileSize / 1024) // XXX: block size 1024 bytes
	attr.Atime = n.sk_file.ModTime
	attr.Mtime = n.sk_file.ModTime
	attr.Ctime = n.sk_file.ModTime
	attr.Crtime = n.sk_file.ModTime
	attr.Mode = os.ModeDir | 0755
	attr.Nlink = 0

	return nil
}

var _ fs.NodeRequestLookuper = (*Dir)(nil)

func (n *Dir) Lookup(ctx context.Context, req *fuse.LookupRequest, resp *fuse.LookupResponse) (fs.Node, error) {
	path := req.Name
	logger.Debugf("req.Name= " + req.Name)
	gd_file, err := n.fs.gd.GetFile(path)
	if err != nil {
		log.Fatalf("Lookup GetFile failed: %v", err)
		return nil, fuse.ENOENT
	}
	if gd_file.IsFolder() {
		return Dir{
			sk_file: gd_file,
			fs:      n.fs,
		}, nil
	} else {
		return File{
			sk_file: gd_file,
			fs:      n.fs,
		}, nil
	}
}

var _ fs.NodeOpener = (*File)(nil)

func (n File) Open(ctx context.Context, req *fuse.OpenRequest, resp *fuse.OpenResponse) (fs.Handle, error) {
	resp.Flags |= fuse.OpenNonSeekable
	return &FileHandle{
		sk_file: n.sk_file,
	}, nil
}

var _ fs.NodeOpener = (*Dir)(nil)

func (n Dir) Open(ctx context.Context, req *fuse.OpenRequest, resp *fuse.OpenResponse) (fs.Handle, error) {
	logger.Debugf("Dir.Open()")
	resp.Flags |= fuse.OpenNonSeekable
	return &DirHandle{
		sk_file: n.sk_file,
	}, nil
}

var _ fs.Handle = (*FileHandle)(nil)

type FileHandle struct {
	sk_file *gdrive.File
}

var _ fs.HandleReleaser = (*FileHandle)(nil)

func (fh *FileHandle) Release(ctx context.Context, req *fuse.ReleaseRequest) error {
	return nil
}

var _ fs.HandleReader = (*FileHandle)(nil)

func (fh *FileHandle) Read(ctx context.Context, req *fuse.ReadRequest, resp *fuse.ReadResponse) error {
	buf := make([]byte, req.Size)
	reader, err := gd.GetFileContents(fh.sk_file)
	if err != nil {
		log.Fatalf("FileHandle Read GetFileContents failed: %v", err)
	}
	n, err := reader.Read(buf)
	if err != nil {
		log.Fatalf("Filehandle Read reader.Read failed: %v", err)
	}
	resp.Data = buf[:n]
	return err
}
<commit_msg>Fix# global `gd` nil pointer dereference, because it is redeclared at this line, this changes the scope<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"golang.org/x/net/context"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"

	skicka "github.com/pellaeon/skicka"
	"github.com/pellaeon/skicka/gdrive"

	"github.com/pellaeon/goas/v3/logger"
)

///////////////////////////////////////////////////////////////////////////
// Global Variables

const BLKSIZE = 4096

var progName = filepath.Base(os.Args[0])

type debugging bool

var (
	gd *gdrive.GDrive

	// The key is only set if encryption is needed (i.e. if -encrypt is
	// provided for an upload, or if an encrypted file is encountered
	// during 'download' or 'cat').
	key []byte

	debug   debugging
	verbose debugging
	quiet   bool

	// Configuration read in from the skicka config file.
	config struct {
		Google struct {
			ClientId     string
			ClientSecret string
			// If set, is appended to all http requests via ?key=XXX.
			ApiKey string
		}
		Encryption struct {
			Salt             string
			Passphrase_hash  string
			Encrypted_key    string
			Encrypted_key_iv string
		}
		Upload struct {
			Ignored_Regexp         []string
			Bytes_per_second_limit int
		}
		Download struct {
			Bytes_per_second_limit int
		}
	}

	// Various statistics gathered along the way. These all should be
	// updated using atomic operations since we often have multiple threads
	// working concurrently for uploads and downloads.
	stats struct {
		DiskReadBytes     int64
		DiskWriteBytes    int64
		UploadBytes       int64
		DownloadBytes     int64
		LocalFilesUpdated int64
		DriveFilesUpdated int64
	}

	// Smaller files will be handled with multiple threads going at once;
	// doing so improves bandwidth utilization since round-trips to the
	// Drive APIs take a while.  (However, we don't want too have too many
	// workers; this would both lead to lots of 403 rate limit errors...)
	nWorkers int
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", progName)
	fmt.Fprintf(os.Stderr, "  %s MOUNTPOINT\n", progName)
	flag.PrintDefaults()
}

func userHomeDir() string {
	if runtime.GOOS == "windows" {
		home := os.Getenv("HOMEDRIVE") + os.Getenv("HOMEPATH")
		if home == "" {
			home = os.Getenv("USERPROFILE")
		}
		return home
	}
	return os.Getenv("HOME")
}

func main() {
	l := logger.NewStandardTimeLogger(os.Stdout, time.StampMicro)
	logger.SetLogger(l)
	logger.SetLevel(logger.LevelDebug)
	log.SetFlags(0)
	log.SetPrefix(progName + ": ")

	// Initialize skicka
	home := userHomeDir()
	tokenCacheFilename := flag.String("tokencache",
		filepath.Join(home, ".skicka.tokencache.json"),
		"OAuth2 token cache file")
	configFilename := flag.String("config",
		filepath.Join(home, ".skicka.config"),
		"Configuration file")
	metadataCacheFilename := flag.String("metadata-cache-file",
		filepath.Join(home, "/.skicka.metadata.cache"),
		"Filename for local cache of Google Drive file metadata")
	nw := flag.Int("num-threads", 4, "Number of threads to use for uploads/downloads")
	vb := flag.Bool("verbose", false, "Enable verbose output")
	dbg := flag.Bool("debug", false, "Enable debugging output")
	qt := flag.Bool("quiet", false, "Suppress non-error messages")
	dumpHTTP := flag.Bool("dump-http", false, "Dump http traffic")
	flakyHTTP := flag.Bool("flaky-http", false, "Add flakiness to http traffic")
	noBrowserAuth := flag.Bool("no-browser-auth", false,
		"Don't try launching browser for authorization")
	flag.Usage = usage
	flag.Parse()

	var mountpoint string
	if flag.NArg() == 1 {
		mountpoint = flag.Arg(0)
	} else {
		mountpoint = *flag.String("mountpoint", "", "Mountpoint")
	}
	if mountpoint == "" {
		usage()
		os.Exit(2)
	}

	nWorkers = *nw

	debug = debugging(*dbg)
	verbose = debugging(*vb || bool(debug))
	quiet = *qt

	skicka.ReadConfigFile(*configFilename)

	// Set up the basic http.Transport.
	transport := http.DefaultTransport
	if tr, ok := transport.(*http.Transport); ok {
		// Increase the default number of open connections per destination host
		// to be enough for the number of goroutines we run concurrently for
		// uploads/downloads; this gives some benefit especially for uploading
		// small files.
		tr.MaxIdleConnsPerHost = 4
	} else {
		skicka.PrintErrorAndExit(fmt.Errorf("DefaultTransport not an *http.Transport?"))
	}
	if *flakyHTTP {
		transport = skicka.NewFlakyTransport(transport)
	}
	if *dumpHTTP {
		//transport = skicka.LoggingTransport{transport: transport}
		// TODO
	}

	// And now upgrade to the OAuth Transport *http.Client.
	client, err := skicka.GetOAuthClient(*tokenCacheFilename, !*noBrowserAuth,
		transport)
	if err != nil {
		skicka.PrintErrorAndExit(fmt.Errorf("error with OAuth2 Authorization: %v ", err))
	}

	// Choose the appropriate callback function for the GDrive object to
	// use for debugging output.
	var dpf func(s string, args ...interface{})
	if debug {
		dpf = skicka.DebugPrint
	} else {
		dpf = skicka.DebugNoPrint
	}

	gd, err = gdrive.New(config.Upload.Bytes_per_second_limit,
		config.Download.Bytes_per_second_limit, dpf, client,
		*metadataCacheFilename, quiet)
	if err != nil {
		log.Fatal(fmt.Errorf("error creating Google Drive "+
			"client: %v", err))
		os.Exit(3)
	}

	c, err := fuse.Mount(mountpoint)
	if err != nil {
		log.Fatalf("Mount failed: %v", err)
	}
	defer c.Close()

	filesys := &FS{
		gd: gd,
	}
	if err := fs.Serve(c, filesys); err != nil {
		log.Fatalf("Serve failed: %v", err)
	}

	// check if the mount process has an error to report
	<-c.Ready
	if err := c.MountError; err != nil {
		log.Fatalf("mount process error: %v", err)
	}
}

type FS struct {
	gd *gdrive.GDrive
}

var _ fs.FS = (*FS)(nil)

func (f *FS) Root() (fs.Node, error) {
	gd_file, err := f.gd.GetFile("/")
	if err != nil {
		logger.Errorf("FS Root() : %v", err)
	}
	return Dir{
		sk_file: gd_file,
		fs:      f,
	}, err
}

var _ fs.FSStatfser = (*FS)(nil)

func (fs *FS) Statfs(ctx context.Context, req *fuse.StatfsRequest, resp *fuse.StatfsResponse) error {
	usage, err := fs.gd.GetDriveUsage()
	if err != nil {
		logger.Errorf("Statfs() error: %v", err)
	}
	resp.Blocks = uint64(usage.Capacity / BLKSIZE)
	resp.Bfree = uint64((usage.Capacity - usage.Used) / BLKSIZE)
	resp.Bavail = uint64((usage.Capacity - usage.Used) / BLKSIZE)
	resp.Files = 9999 //XXX
	resp.Ffree = 9999 //XXX
	resp.Bsize = BLKSIZE
	resp.Namelen = 32767 //http://www.aurelp.com/2014/09/10/what-is-the-maximum-name-length-for-a-file-on-google-drive/
	resp.Frsize = BLKSIZE
	req.Respond(resp)
	return err
}

var _ fs.Node = (*Dir)(nil)

type Dir struct {
	sk_file *gdrive.File
	fs      *FS
}

func (n Dir) Attr(ctx context.Context, attr *fuse.Attr) error {
	var sum uint64
	for _, c := range n.sk_file.Id {
		sum += uint64(c)
	}
	attr.Inode = sum
	attr.Size = 4096 // XXX
	attr.Blocks = 0
	attr.Atime = n.sk_file.ModTime
	attr.Mtime = n.sk_file.ModTime
	attr.Ctime = n.sk_file.ModTime
	attr.Crtime = n.sk_file.ModTime
	attr.Mode = os.ModeDir | 0755
	attr.Nlink = 0

	return nil
}

type DirHandle struct {
	sk_file *gdrive.File
}

var _ fs.HandleReadDirAller = (*DirHandle)(nil)

func (dh *DirHandle) ReadDirAll(ctx context.Context) ([]fuse.Dirent, error) {
	logger.Debugf("ReadDirAll %s", dh.sk_file.Path)
	files_in_folder, err := gd.GetFilesInFolder(dh.sk_file.Path)
	logger.Debugf("ReadDirAll %+v", files_in_folder)
	if err != nil {
		log.Fatalf("ReadDirAll failed: %v", err)
	}
	var res []fuse.Dirent
	for _, file := range files_in_folder {
		var de fuse.Dirent
		de.Name = file.Path[strings.LastIndex(file.Path, "/")+1:]
		if file.IsFolder() {
			de.Type = fuse.DT_Dir
		} else {
			de.Type = fuse.DT_File
		}
		sum := uint64(0)
		for _, c := range file.Id {
			sum += uint64(c)
		}
		de.Inode = sum
		res = append(res, de)
	}
	return res, err
}

var _ fs.Node = (*File)(nil)

type File struct {
	sk_file *gdrive.File
	fs      *FS
}

func (n File) Attr(ctx context.Context, attr *fuse.Attr) error {
	sum := uint64(0)
	for _, c := range n.sk_file.Id {
		sum += uint64(c)
	}
	attr.Inode = sum
	attr.Size = uint64(n.sk_file.FileSize)
	attr.Blocks = uint64(n.sk_file.FileSize / 1024) // XXX: block size 1024 bytes
	attr.Atime = n.sk_file.ModTime
	attr.Mtime = n.sk_file.ModTime
	attr.Ctime = n.sk_file.ModTime
	attr.Crtime = n.sk_file.ModTime
	attr.Mode = os.ModeDir | 0755
	attr.Nlink = 0

	return nil
}

var _ fs.NodeRequestLookuper = (*Dir)(nil)

func (n *Dir) Lookup(ctx context.Context, req *fuse.LookupRequest, resp *fuse.LookupResponse) (fs.Node, error) {
	path := req.Name
	logger.Debugf("req.Name= " + req.Name)
	gd_file, err := n.fs.gd.GetFile(path)
	if err != nil {
		log.Fatalf("Lookup GetFile failed: %v", err)
		return nil, fuse.ENOENT
	}
	if gd_file.IsFolder() {
		return Dir{
			sk_file: gd_file,
			fs:      n.fs,
		}, nil
	} else {
		return File{
			sk_file: gd_file,
			fs:      n.fs,
		}, nil
	}
}

var _ fs.NodeOpener = (*File)(nil)

func (n File) Open(ctx context.Context, req *fuse.OpenRequest, resp *fuse.OpenResponse) (fs.Handle, error) {
	resp.Flags |= fuse.OpenNonSeekable
	return &FileHandle{
		sk_file: n.sk_file,
	}, nil
}

var _ fs.NodeOpener = (*Dir)(nil)

func (n Dir) Open(ctx context.Context, req *fuse.OpenRequest, resp *fuse.OpenResponse) (fs.Handle, error) {
	logger.Debugf("Dir.Open()")
	resp.Flags |= fuse.OpenNonSeekable
	return &DirHandle{
		sk_file: n.sk_file,
	}, nil
}

var _ fs.Handle = (*FileHandle)(nil)

type FileHandle struct {
	sk_file *gdrive.File
}

var _ fs.HandleReleaser = (*FileHandle)(nil)

func (fh *FileHandle) Release(ctx context.Context, req *fuse.ReleaseRequest) error {
	return nil
}

var _ fs.HandleReader = (*FileHandle)(nil)

func (fh *FileHandle) Read(ctx context.Context, req *fuse.ReadRequest, resp *fuse.ReadResponse) error {
	buf := make([]byte, req.Size)
	reader, err := gd.GetFileContents(fh.sk_file)
	if err != nil {
		log.Fatalf("FileHandle Read GetFileContents failed: %v", err)
	}
	n, err := reader.Read(buf)
	if err != nil {
		log.Fatalf("Filehandle Read reader.Read failed: %v", err)
	}
	resp.Data = buf[:n]
	return err
}
<|endoftext|>