<commit_before>package commands

import (
	"bufio"
	"io/ioutil"
	"os"
	"strings"

	"github.com/github/git-lfs/lfs"
	"github.com/github/git-lfs/vendor/_nuts/github.com/spf13/cobra"
)

var (
	untrackCmd = &cobra.Command{
		Use:   "untrack",
		Short: "Remove an entry from .gitattributes",
		Run:   untrackCommand,
	}
)

// untrackCommand takes a list of paths as an argument, and removes each path from the
// default attribtues file (.gitattributes), if it exists.
func untrackCommand(cmd *cobra.Command, args []string) {
	lfs.InstallHooks(false)

	if len(args) < 1 {
		Print("git lfs untrack <path> [path]*")
		return
	}

	data, err := ioutil.ReadFile(".gitattributes")
	if err != nil {
		return
	}

	attributes := strings.NewReader(string(data))

	attributesFile, err := os.Create(".gitattributes")
	if err != nil {
		Print("Error opening .gitattributes for writing")
		return
	}

	scanner := bufio.NewScanner(attributes)

	// Iterate through each line of the attributes file and rewrite it,
	// if the path was meant to be untracked, omit it, and print a message instead.
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, "filter=lfs") {
			fields := strings.Fields(line)
			removeThisPath := false
			for _, t := range args {
				if t == fields[0] {
					removeThisPath = true
				}
			}

			if !removeThisPath {
				attributesFile.WriteString(line + "\n")
			} else {
				Print("Untracking %s", fields[0])
			}
		}
	}

	attributesFile.Close()
}

func init() {
	RootCmd.AddCommand(untrackCmd)
}
<commit_msg>ラララララ ラー ウウウ フフフ<commit_after>package commands

import (
	"bufio"
	"io/ioutil"
	"os"
	"strings"

	"github.com/github/git-lfs/lfs"
	"github.com/github/git-lfs/vendor/_nuts/github.com/spf13/cobra"
)

var (
	untrackCmd = &cobra.Command{
		Use:   "untrack",
		Short: "Remove an entry from .gitattributes",
		Run:   untrackCommand,
	}
)

// untrackCommand takes a list of paths as an argument, and removes each path from the
// default attribtues file (.gitattributes), if it exists.
func untrackCommand(cmd *cobra.Command, args []string) {
	if lfs.LocalGitDir == "" {
		Print("Not a git repository.")
		os.Exit(128)
	}
	if lfs.LocalWorkingDir == "" {
		Print("This operation must be run in a work tree.")
		os.Exit(128)
	}

	lfs.InstallHooks(false)

	if len(args) < 1 {
		Print("git lfs untrack <path> [path]*")
		return
	}

	data, err := ioutil.ReadFile(".gitattributes")
	if err != nil {
		return
	}

	attributes := strings.NewReader(string(data))

	attributesFile, err := os.Create(".gitattributes")
	if err != nil {
		Print("Error opening .gitattributes for writing")
		return
	}

	scanner := bufio.NewScanner(attributes)

	// Iterate through each line of the attributes file and rewrite it,
	// if the path was meant to be untracked, omit it, and print a message instead.
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, "filter=lfs") {
			fields := strings.Fields(line)
			removeThisPath := false
			for _, t := range args {
				if t == fields[0] {
					removeThisPath = true
				}
			}

			if !removeThisPath {
				attributesFile.WriteString(line + "\n")
			} else {
				Print("Untracking %s", fields[0])
			}
		}
	}

	attributesFile.Close()
}

func init() {
	RootCmd.AddCommand(untrackCmd)
}
<|endoftext|>


<commit_before>package compilers

import (
	"bufio"
	"fmt"
	"github.com/BurntSushi/toml"
	"github.com/albrow/scribble/config"
	"github.com/albrow/scribble/context"
	"github.com/albrow/scribble/util"
	"github.com/howeyc/fsnotify"
	"github.com/russross/blackfriday"
	"github.com/wsxiaoys/terminal/color"
	"html/template"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

// PostsCompilerType represents a type capable of compiling post files.
type PostsCompilerType struct {
	pathMatch string
	// createdDirs keeps track of the directories that were created in config.DestDir.
	// It is used in the RemoveOld method.
	createdDirs []string
}

// PostCompiler is an instatiation of PostCompilerType
var PostsCompiler = PostsCompilerType{
	pathMatch: "",
}

// Post is an in-memory representation of the metadata for a given post.
// Much of this data comes from the toml frontmatter.
type Post struct {
	Title       string    `toml:"title"`
	Author      string    `toml:"author"`
	Description string    `toml:"description"`
	Date        time.Time `toml:"date"`
	// the url for the post, not including protocol or domain name (useful for creating links)
	Url template.URL `toml:"-"`
	// the html content for the post (parsed from markdown source)
	Content template.HTML `toml:"-"`
	// the full source path
	src string `toml:"-"`
	// the layout tmpl file to be used for the post
	Layout string `toml:"layout"`
	// the html template that can be used to render the post
	template *template.Template
}

var (
	// a slice of all posts
	posts = []*Post{}
	// a map of source path to post
	postsMap = map[string]*Post{}
)

// Init should be called before any other methods. In this case, Init
// sets up the pathMatch variable based on config.SourceDir and config.PostsDir
// and adds the Posts helper function to FuncMap.
func (p *PostsCompilerType) Init() {
	p.pathMatch = filepath.Join(config.PostsDir, "*.md")
	// Add the posts function to FuncMap
	context.FuncMap["Posts"] = Posts
}

// CompileMatchFunc returns a MatchFunc which will return true for
// any files which match a given pattern. In this case, the pattern
// is any file that is inside config.PostsDir and ends in ".md", excluding
// hidden files and directories (which start with a ".") but not those
// which start with an underscore.
func (p *PostsCompilerType) CompileMatchFunc() MatchFunc {
	return pathMatchFunc(p.pathMatch, true, false)
}

// WatchMatchFunc returns a MatchFunc which will return true for
// any files which match a given pattern. In this case, the pattern
// is the same as it is for CompileMatchFunc.
func (p *PostsCompilerType) WatchMatchFunc() MatchFunc {
	// PostsCompiler needs to watch all posts in the posts dir,
	// but also needs to watch all the *tmpl files in the posts
	// layouts dir, because if those change, it affects the way
	// posts are rendered.
	postsMatch := pathMatchFunc(p.pathMatch, true, false)
	postLayoutsMatch := pathMatchFunc(filepath.Join(config.PostLayoutsDir, "*.tmpl"), true, false)
	// unionMatchFuncs combines these two cases and returns a MatchFunc
	// which will return true if either matches. This allows us to watch
	// for changes in both the posts dir and the posts layouts dir.
	return unionMatchFuncs(postsMatch, postLayoutsMatch)
}

// Compile compiles the file at srcPath. The caller will only
// call this function for files which belong to PostsCompiler
// according to the MatchFunc. Behavior for any other file is
// undefined. Compile will output the compiled result to the appropriate
// location in config.DestDir.
func (p *PostsCompilerType) Compile(srcPath string) error {
	// Get the parsed post object and determine dest path
	post := getOrCreatePostFromPath(srcPath)
	srcFilename := filepath.Base(srcPath)
	destPath := fmt.Sprintf("%s/%s", config.DestDir, strings.TrimSuffix(srcFilename, ".md"))
	destIndexFilePath := filepath.Join(destPath, "index.html")
	color.Printf("@g    CREATE: %s -> %s\n", srcPath, destIndexFilePath)

	// Create the index file
	destFile, err := util.CreateFileWithPath(destIndexFilePath)
	if err != nil {
		return err
	}

	// Parse content and frontmatter, then set the appropriate layout based on
	// the layout key in the frontmatter
	if err := post.parse(); err != nil {
		return err
	}

	// Write to destFile by executing the template
	postContext := context.CopyContext()
	postContext["Post"] = post
	if err := post.template.Execute(destFile, postContext); err != nil {
		return fmt.Errorf("ERROR compiling html template for posts: %s", err.Error())
	}

	// Add the created dir to the list of created dirs
	p.createdDirs = append(p.createdDirs, destPath)

	return nil
}

// CompileAll compiles zero or more files identified by srcPaths.
// It works simply by calling Compile for each path. The caller is
// responsible for only passing in files that belong to AceCompiler
// according to the MatchFunc. Behavior for any other file is undefined.
func (p *PostsCompilerType) CompileAll(srcPaths []string) error {
	fmt.Println("--> compiling posts")
	for _, srcPath := range srcPaths {
		if err := p.Compile(srcPath); err != nil {
			return err
		}
	}
	return nil
}

func (p *PostsCompilerType) FileChanged(srcPath string, ev fsnotify.FileEvent) error {
	// Because of the way we set up the watcher, there are two possible
	// cases here.
	// 1) A template in the post layouts dir was changed. In this case,
	// we would ideally recompile all the posts that used that layout.
	// 2) A markdown file corresponding to a single post was changed. In this
	// case, ideally we only recompile the post that was changed. We need to
	// take into account any rename, delete, or create events and how they
	// affect the output files in destDir.
	switch filepath.Ext(srcPath) {
	case ".md":
		// TODO: Be more intelligent here? If a single post file was midified,
		// we can simply recompile that post. We would also need to take into
		// account the subtle differences between rename, create, and delete
		// events. For now, recompile all posts.
		if err := recompileAllForCompiler(p); err != nil {
			return err
		}
		return nil
	case ".tmpl":
		// TODO: Analyze post files and be more intelligent here?
		// When a post layout changes, only recompile the posts that
		// use that layout. For now, recompile all posts.
		if err := recompileAllForCompiler(p); err != nil {
			return err
		}
		return nil
	}
	return nil
}

func (p *PostsCompilerType) RemoveOld() error {
	// Simply iterate through createdDirs and remove each of them
	// NOTE: this is different from the other compilers because each
	// post gets created as an index.html file inside some directory
	// (for prettier urls). So instead of removing files, we're removing
	// directories.
	for _, dir := range p.createdDirs {
		if err := util.RemoveAllIfExists(dir); err != nil {
			return err
		}
	}
	return nil
}

// Posts returns up to limit posts, sorted by date. If limit is 0,
// it returns all posts. If limit is greater than len(posts), it returns
// all posts.
func Posts(limit ...int) []*Post {
	// Sort the posts by date
	sortedPosts := make([]*Post, len(posts))
	copy(sortedPosts, posts)
	sort.Sort(PostsByDate(sortedPosts))

	// Return up to limit posts
	if len(limit) == 0 || limit[0] == 0 || limit[0] > len(sortedPosts) {
		return sortedPosts
	} else {
		return sortedPosts[:limit[0]]
	}
}

func createPostFromPath(path string) *Post {
	// create post object
	name := strings.TrimSuffix(filepath.Base(path), filepath.Ext(path))
	p := &Post{
		Url: template.URL("/" + name),
		src: path,
	}
	posts = append(posts, p)
	postsMap[path] = p
	return p
}

func getPostByPath(path string) *Post {
	return postsMap[path]
}

func getOrCreatePostFromPath(path string) *Post {
	if p, found := postsMap[path]; found {
		return p
	} else {
		return createPostFromPath(path)
	}
}

// parse reads from the source file and sets the content and metadata fields
// for the post. It also creates a new template for the post using the layout
// field of the frontmatter.
func (p *Post) parse() error {
	// Open the source file
	file, err := os.Open(p.src)
	if err != nil {
		return err
	}
	r := bufio.NewReader(file)

	// Split the file into frontmatter and markdown content
	frontMatter, content, err := util.SplitFrontMatter(r)
	if err != nil {
		return err
	}

	// Decode the frontmatter
	if _, err := toml.Decode(frontMatter, p); err != nil {
		return err
	}

	// Parse the markdown content and set p.Content
	p.Content = template.HTML(blackfriday.MarkdownCommon([]byte(content)))

	// Create a template for the post
	if p.Layout == "" {
		return fmt.Errorf("Could not find layout definition in toml frontmatter for post: %s", p.src)
	}

	// When we call template.ParseFiles, we want the post layout file to be first,
	// so it is the one that will be executed.
	postLayoutFile := filepath.Join(config.PostLayoutsDir, p.Layout)
	otherLayoutFiles, err := filepath.Glob(filepath.Join(config.LayoutsDir, "*.tmpl"))
	if err != nil {
		return err
	}
	allFiles := append([]string{postLayoutFile}, otherLayoutFiles...)
	// Parse the includes files if there are any
	if config.IncludesDir != "" {
		includeFiles, err := filepath.Glob(filepath.Join(config.IncludesDir, "*tmpl"))
		if err != nil {
			return err
		}
		allFiles = append(allFiles, includeFiles...)
	}
	tmpl, err := template.ParseFiles(allFiles...)
	if err != nil {
		return err
	}
	p.template = tmpl

	return nil
}

// The PostsByDate type is used only for sorting
type PostsByDate []*Post

func (p PostsByDate) Len() int {
	return len(p)
}

func (p PostsByDate) Less(i, j int) bool {
	return p[i].Date.Before(p[j].Date)
}

func (p PostsByDate) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}
<commit_msg>Watch includes and layouts dir and recompile posts when they change.<commit_after>package compilers

import (
	"bufio"
	"fmt"
	"github.com/BurntSushi/toml"
	"github.com/albrow/scribble/config"
	"github.com/albrow/scribble/context"
	"github.com/albrow/scribble/util"
	"github.com/howeyc/fsnotify"
	"github.com/russross/blackfriday"
	"github.com/wsxiaoys/terminal/color"
	"html/template"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

// PostsCompilerType represents a type capable of compiling post files.
type PostsCompilerType struct {
	pathMatch string
	// createdDirs keeps track of the directories that were created in config.DestDir.
	// It is used in the RemoveOld method.
	createdDirs []string
}

// PostCompiler is an instatiation of PostCompilerType
var PostsCompiler = PostsCompilerType{
	pathMatch: "",
}

// Post is an in-memory representation of the metadata for a given post.
// Much of this data comes from the toml frontmatter.
type Post struct {
	Title       string    `toml:"title"`
	Author      string    `toml:"author"`
	Description string    `toml:"description"`
	Date        time.Time `toml:"date"`
	// the url for the post, not including protocol or domain name (useful for creating links)
	Url template.URL `toml:"-"`
	// the html content for the post (parsed from markdown source)
	Content template.HTML `toml:"-"`
	// the full source path
	src string `toml:"-"`
	// the layout tmpl file to be used for the post
	Layout string `toml:"layout"`
	// the html template that can be used to render the post
	template *template.Template
}

var (
	// a slice of all posts
	posts = []*Post{}
	// a map of source path to post
	postsMap = map[string]*Post{}
)

// Init should be called before any other methods. In this case, Init
// sets up the pathMatch variable based on config.SourceDir and config.PostsDir
// and adds the Posts helper function to FuncMap.
func (p *PostsCompilerType) Init() {
	p.pathMatch = filepath.Join(config.PostsDir, "*.md")
	// Add the posts function to FuncMap
	context.FuncMap["Posts"] = Posts
}

// CompileMatchFunc returns a MatchFunc which will return true for
// any files which match a given pattern. In this case, the pattern
// is any file that is inside config.PostsDir and ends in ".md", excluding
// hidden files and directories (which start with a ".") but not those
// which start with an underscore.
func (p *PostsCompilerType) CompileMatchFunc() MatchFunc {
	return pathMatchFunc(p.pathMatch, true, false)
}

// WatchMatchFunc returns a MatchFunc which will return true for
// any files which match a given pattern. In this case, the pattern
// is the same as it is for CompileMatchFunc.
func (p *PostsCompilerType) WatchMatchFunc() MatchFunc {
	// PostsCompiler needs to watch all posts in the posts dir,
	// but also needs to watch all the *tmpl files in the posts
	// layouts dir, layouts dir, and includes dir. Because if those
	// change, it may affect the way posts are rendered.
	postsMatch := pathMatchFunc(p.pathMatch, true, false)
	layoutsMatch := pathMatchFunc(filepath.Join(config.LayoutsDir, "*.tmpl"), true, false)
	postLayoutsMatch := pathMatchFunc(filepath.Join(config.PostLayoutsDir, "*.tmpl"), true, false)
	// unionMatchFuncs combines these two cases and returns a MatchFunc
	// which will return true if either matches. This allows us to watch
	// for changes in both the posts dir and the posts layouts dir.
	allMatch := unionMatchFuncs(postsMatch, layoutsMatch, postLayoutsMatch)
	if config.IncludesDir != "" {
		// We also want to watch includes if there are any
		includesMatch := pathMatchFunc(filepath.Join(config.IncludesDir, "*.tmpl"), true, false)
		allMatch = unionMatchFuncs(allMatch, includesMatch)
	}
	return allMatch
}

// Compile compiles the file at srcPath. The caller will only
// call this function for files which belong to PostsCompiler
// according to the MatchFunc. Behavior for any other file is
// undefined. Compile will output the compiled result to the appropriate
// location in config.DestDir.
func (p *PostsCompilerType) Compile(srcPath string) error {
	// Get the parsed post object and determine dest path
	post := getOrCreatePostFromPath(srcPath)
	srcFilename := filepath.Base(srcPath)
	destPath := fmt.Sprintf("%s/%s", config.DestDir, strings.TrimSuffix(srcFilename, ".md"))
	destIndexFilePath := filepath.Join(destPath, "index.html")
	color.Printf("@g    CREATE: %s -> %s\n", srcPath, destIndexFilePath)

	// Create the index file
	destFile, err := util.CreateFileWithPath(destIndexFilePath)
	if err != nil {
		return err
	}

	// Parse content and frontmatter, then set the appropriate layout based on
	// the layout key in the frontmatter
	if err := post.parse(); err != nil {
		return err
	}

	// Write to destFile by executing the template
	postContext := context.CopyContext()
	postContext["Post"] = post
	if err := post.template.Execute(destFile, postContext); err != nil {
		return fmt.Errorf("ERROR compiling html template for posts: %s", err.Error())
	}

	// Add the created dir to the list of created dirs
	p.createdDirs = append(p.createdDirs, destPath)

	return nil
}

// CompileAll compiles zero or more files identified by srcPaths.
// It works simply by calling Compile for each path. The caller is
// responsible for only passing in files that belong to AceCompiler
// according to the MatchFunc. Behavior for any other file is undefined.
func (p *PostsCompilerType) CompileAll(srcPaths []string) error {
	fmt.Println("--> compiling posts")
	for _, srcPath := range srcPaths {
		if err := p.Compile(srcPath); err != nil {
			return err
		}
	}
	return nil
}

func (p *PostsCompilerType) FileChanged(srcPath string, ev fsnotify.FileEvent) error {
	// Because of the way we set up the watcher, there are two possible
	// cases here.
	// 1) A template in the post layouts dir was changed. In this case,
	// we would ideally recompile all the posts that used that layout.
	// 2) A markdown file corresponding to a single post was changed. In this
	// case, ideally we only recompile the post that was changed. We need to
	// take into account any rename, delete, or create events and how they
	// affect the output files in destDir.
	switch filepath.Ext(srcPath) {
	case ".md":
		// TODO: Be more intelligent here? If a single post file was midified,
		// we can simply recompile that post. We would also need to take into
		// account the subtle differences between rename, create, and delete
		// events. For now, recompile all posts.
		if err := recompileAllForCompiler(p); err != nil {
			return err
		}
		return nil
	case ".tmpl":
		// TODO: Analyze post files and be more intelligent here?
		// When a post layout changes, only recompile the posts that
		// use that layout. For now, recompile all posts.
		if err := recompileAllForCompiler(p); err != nil {
			return err
		}
		return nil
	}
	return nil
}

func (p *PostsCompilerType) RemoveOld() error {
	// Simply iterate through createdDirs and remove each of them
	// NOTE: this is different from the other compilers because each
	// post gets created as an index.html file inside some directory
	// (for prettier urls). So instead of removing files, we're removing
	// directories.
	for _, dir := range p.createdDirs {
		if err := util.RemoveAllIfExists(dir); err != nil {
			return err
		}
	}
	return nil
}

// Posts returns up to limit posts, sorted by date. If limit is 0,
// it returns all posts. If limit is greater than len(posts), it returns
// all posts.
func Posts(limit ...int) []*Post {
	// Sort the posts by date
	sortedPosts := make([]*Post, len(posts))
	copy(sortedPosts, posts)
	sort.Sort(PostsByDate(sortedPosts))

	// Return up to limit posts
	if len(limit) == 0 || limit[0] == 0 || limit[0] > len(sortedPosts) {
		return sortedPosts
	} else {
		return sortedPosts[:limit[0]]
	}
}

func createPostFromPath(path string) *Post {
	// create post object
	name := strings.TrimSuffix(filepath.Base(path), filepath.Ext(path))
	p := &Post{
		Url: template.URL("/" + name),
		src: path,
	}
	posts = append(posts, p)
	postsMap[path] = p
	return p
}

func getPostByPath(path string) *Post {
	return postsMap[path]
}

func getOrCreatePostFromPath(path string) *Post {
	if p, found := postsMap[path]; found {
		return p
	} else {
		return createPostFromPath(path)
	}
}

// parse reads from the source file and sets the content and metadata fields
// for the post. It also creates a new template for the post using the layout
// field of the frontmatter.
func (p *Post) parse() error {
	// Open the source file
	file, err := os.Open(p.src)
	if err != nil {
		return err
	}
	r := bufio.NewReader(file)

	// Split the file into frontmatter and markdown content
	frontMatter, content, err := util.SplitFrontMatter(r)
	if err != nil {
		return err
	}

	// Decode the frontmatter
	if _, err := toml.Decode(frontMatter, p); err != nil {
		return err
	}

	// Parse the markdown content and set p.Content
	p.Content = template.HTML(blackfriday.MarkdownCommon([]byte(content)))

	// Create a template for the post
	if p.Layout == "" {
		return fmt.Errorf("Could not find layout definition in toml frontmatter for post: %s", p.src)
	}

	// When we call template.ParseFiles, we want the post layout file to be first,
	// so it is the one that will be executed.
	postLayoutFile := filepath.Join(config.PostLayoutsDir, p.Layout)
	otherLayoutFiles, err := filepath.Glob(filepath.Join(config.LayoutsDir, "*.tmpl"))
	if err != nil {
		return err
	}
	allFiles := append([]string{postLayoutFile}, otherLayoutFiles...)
	// Parse the includes files if there are any
	if config.IncludesDir != "" {
		includeFiles, err := filepath.Glob(filepath.Join(config.IncludesDir, "*tmpl"))
		if err != nil {
			return err
		}
		allFiles = append(allFiles, includeFiles...)
	}
	tmpl, err := template.ParseFiles(allFiles...)
	if err != nil {
		return err
	}
	p.template = tmpl

	return nil
}

// The PostsByDate type is used only for sorting
type PostsByDate []*Post

func (p PostsByDate) Len() int {
	return len(p)
}

func (p PostsByDate) Less(i, j int) bool {
	return p[i].Date.Before(p[j].Date)
}

func (p PostsByDate) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
}
<|endoftext|>


<commit_before>package raft

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

	"github.com/HelloCodeMing/raft-rocks/pb"
	"github.com/HelloCodeMing/raft-rocks/store"
	"github.com/HelloCodeMing/raft-rocks/utils"
	"github.com/golang/glog"
)

// Hold state of raft, coordinate the updating and reading
type raftState struct {
	sync.RWMutex
	log       *store.LogStorage
	persister store.Persister
	applyCh   chan<- *ApplyMsg

	// once initialized, they will not be changed
	Me       int
	NumPeers int

	Role        pb.RaftRole
	CurrentTerm int32
	VotedFor    int32

	// maintained by leader, record follower's repliate and commit state
	MatchIndex []int
	NextIndex  []int

	// in raft paper, these two field is volatile
	// when restart, we need replay all WAL to recover state
	// otherwise, we have to do checkpoint, and record where has beed commited and applied to state machine
	// but when it comes to RocksDB like storage engine, it has its own WAL, so the storage itself is durable
	// but disappointing, even if reading RocksDB's WAL to get LastApplied is possible, it's too tricky
	// An elegant way to avoid 'double WAL' may be stop use WAL in RocksDB, but only in Raft,
	// and in raft we could perioidly do checkpoint/snapshot, persist LastApplied, as a result, we just need to
	// restore state machine from checkpoint, and replay WAL.
	//
	// All above is just my imagination, the real implementation is too naive. When LastApplied is updated, store it in persister.
	CommitIndex int // index of highest log entry known to be committed
	LastApplied int // index of highest log entry applied to state machine

	readLease time.Time
}

func (s *raftState) getRole() pb.RaftRole {
	s.RLock()
	defer s.RUnlock()
	return s.Role
}

// maybe we could cache this field in Raft, to avoid lock
func (s *raftState) getTerm() int32 {
	s.RLock()
	defer s.RUnlock()
	return s.CurrentTerm
}

func (s *raftState) getCommited() int {
	s.RLock()
	defer s.RUnlock()
	return s.CommitIndex
}

const (
	currentTerm = "currentTerm"
	votedFor    = "votedFor"
	lastApplied = "lastApplied"
)

func (s *raftState) persist() {
	s.persister.StoreInt32(currentTerm, s.CurrentTerm)
	s.persister.StoreInt32(votedFor, s.VotedFor)
}

func (s *raftState) restore() {
	var ok bool
	s.CurrentTerm, ok = s.persister.LoadInt32(currentTerm)
	s.VotedFor, ok = s.persister.LoadInt32(votedFor)
	if !ok {
		s.VotedFor = -1
	}
	apply, ok := s.persister.LoadInt32(lastApplied)
	if ok {
		s.LastApplied = int(apply)
	}
}

func (s *raftState) commitUntil(index int) {
	s.CommitIndex = index
}

func (s *raftState) applyOne() {
	s.LastApplied++
	s.persister.StoreInt32(lastApplied, int32(s.LastApplied))
}

func (s *raftState) checkNewTerm(newTerm int32) bool {
	if s.getTerm() < newTerm {
		s.becomeFollower(-1, newTerm)
		return true
	}
	return false
}
func (s *raftState) checkFollowerCommit(leaderCommit int) bool {
	s.Lock()
	defer s.Unlock()
	// Have uncommited log, but leader has commited them
	if s.CommitIndex < leaderCommit && s.CommitIndex < s.log.LastIndex() {
		s.commitUntil(minInt(leaderCommit, int(s.log.LastIndex())))
		s.checkApply()
		return true
	}
	return false
}

// any command to apply
func (s *raftState) checkApply() {
	old := s.LastApplied
	log := s.log
	for s.CommitIndex > s.LastApplied {
		s.applyOne()
		if log.LastIndex() < s.LastApplied {
			panic(s)
		}
		entry := log.At(s.LastApplied)
		if entry == nil {
			break
		}
		msg := &ApplyMsg{Command: entry}
		s.applyCh <- msg
	}
	if s.LastApplied > old {
		glog.V(utils.VDebug).Infof("%s Applyed commands until index=%d", s.String(), s.LastApplied)
	}
}

func (s *raftState) changeRole(role pb.RaftRole) {
	s.Lock()
	defer s.Unlock()
	s.Role = role

}

func (s *raftState) becomeFollowerUnlocked(candidateID int32, term int32) {
	s.CurrentTerm = term
	s.VotedFor = candidateID
	s.Role = pb.RaftRole_Follower
	s.persist()
}

func (s *raftState) becomeFollower(candidateID int32, term int32) {
	s.Lock()
	defer s.Unlock()
	s.CurrentTerm = term
	s.VotedFor = candidateID
	s.Role = pb.RaftRole_Follower
	s.MatchIndex = nil
	s.NextIndex = nil
	s.persist()
}

func (s *raftState) becomeCandidate() int32 {
	s.Lock()
	defer s.Unlock()
	s.Role = pb.RaftRole_Candidate
	s.CurrentTerm++
	s.VotedFor = int32(s.Me)
	s.MatchIndex = nil
	s.NextIndex = nil
	s.persist()
	return s.CurrentTerm
}

func (s *raftState) becomeLeader() {
	s.Lock()
	defer s.Unlock()
	s.Role = pb.RaftRole_Leader
	s.MatchIndex = make([]int, s.NumPeers)
	s.NextIndex = make([]int, s.NumPeers)
	lastIndex := s.log.LastIndex()
	for i := range s.NextIndex {
		s.NextIndex[i] = lastIndex + 1
	}
}

func (s *raftState) replicatedToPeer(peer int, index int) {
	s.Lock()
	defer s.Unlock()
	s.MatchIndex[peer] = maxInt(index, s.MatchIndex[peer])
	s.NextIndex[peer] = s.MatchIndex[peer] + 1
	if s.checkLeaderCommit() {
		s.checkApply()
	}
}

// update  commit index, use matchIndex from peers
// NOTE: It's not thread-safe, should be synchronized by external lock
func (s *raftState) checkLeaderCommit() (updated bool) {
	lowerIndex := s.CommitIndex + 1
	upperIndex := 0
	// find max matchIndex
	for _, x := range s.MatchIndex {
		if x > upperIndex {
			upperIndex = x
		}
	}
	// if N > commitIndex, a majority of match[i] >= N, and log[N].term == currentTerm
	// set commitIndex = N
	for N := upperIndex; N >= lowerIndex && s.log.At(N).Term == s.CurrentTerm; N-- {
		// count match[i] >= N
		cnt := 1
		for i, x := range s.MatchIndex {
			if i != s.Me && x >= N {
				cnt++
			}
		}
		if cnt >= s.majority() {
			s.commitUntil(N)
			glog.V(utils.VDebug).Infof("%s Leader update commitIndex: %d", s.String(), s.CommitIndex)
			updated = true
			break
		}
	}
	return
}

func (s *raftState) retreatForPeer(peer int) int {
	s.Lock()
	defer s.Unlock()
	idx := &s.NextIndex[peer]
	log := s.log
	*idx = minInt(*idx, log.LastIndex())
	oldTerm := log.At(*idx).Term
	for *idx > 0 && log.At(*idx).Term == oldTerm {
		*idx--
	}
	*idx = maxInt(*idx, 1)
	return *idx
}

func (s *raftState) toReplicate(peer int) int {
	return s.NextIndex[peer]
}

// If lease is granted in this term, and later than the old one, extend the lease
func (s *raftState) updateReadLease(term int32, lease time.Time) {
	if s.readLease.Before(time.Now()) && term == s.CurrentTerm && lease.After(s.readLease) {
		s.readLease = lease
		glog.Infof("%s Update read lease to %v", s.String(), lease)
	}
}

func (s *raftState) String() string {
	return fmt.Sprintf("Raft<%d:%d>", s.Me, s.CurrentTerm)
}

func (s *raftState) dump(writer io.Writer) {
	state := map[string]interface{}{
		"meta":         s,
		"LogLastIndex": s.log.LastIndex(),
	}
	buf, err := json.MarshalIndent(state, "", "\t")
	if err != nil {
		panic(err)
	}
	writer.Write(buf)
}

func (s *raftState) majority() int {
	return s.NumPeers/2 + 1
}

func makeRaftState(log *store.LogStorage, persister store.Persister, applyCh chan<- *ApplyMsg, numPeers int, me int) *raftState {
	s := &raftState{
		log:         log,
		applyCh:     applyCh,
		NumPeers:    numPeers,
		Me:          me,
		CurrentTerm: 0,
		VotedFor:    -1,
		Role:        pb.RaftRole_Follower,
		readLease:   time.Now(),
		persister:   persister,
	}
	s.restore()
	glog.Infof("Restore raft state: CurrentTerm: %d, CommitIndex: %d, LastApplied: %d", s.CurrentTerm, s.CommitIndex, s.LastApplied)
	return s
}
<commit_msg>optimize leader commit<commit_after>package raft

import (
	"encoding/json"
	"fmt"
	"io"
	"sort"
	"sync"
	"time"

	"github.com/HelloCodeMing/raft-rocks/pb"
	"github.com/HelloCodeMing/raft-rocks/store"
	"github.com/HelloCodeMing/raft-rocks/utils"
	"github.com/golang/glog"
)

// Hold state of raft, coordinate the updating and reading
type raftState struct {
	sync.RWMutex
	log       *store.LogStorage
	persister store.Persister
	applyCh   chan<- *ApplyMsg

	// once initialized, they will not be changed
	Me       int
	NumPeers int

	Role        pb.RaftRole
	CurrentTerm int32
	VotedFor    int32

	// maintained by leader, record follower's repliate and commit state
	MatchIndex []int
	NextIndex  []int

	// in raft paper, these two field is volatile
	// when restart, we need replay all WAL to recover state
	// otherwise, we have to do checkpoint, and record where has beed commited and applied to state machine
	// but when it comes to RocksDB like storage engine, it has its own WAL, so the storage itself is durable
	// but disappointing, even if reading RocksDB's WAL to get LastApplied is possible, it's too tricky
	// An elegant way to avoid 'double WAL' may be stop use WAL in RocksDB, but only in Raft,
	// and in raft we could perioidly do checkpoint/snapshot, persist LastApplied, as a result, we just need to
	// restore state machine from checkpoint, and replay WAL.
	//
	// All above is just my imagination, the real implementation is too naive. When LastApplied is updated, store it in persister.
	CommitIndex int // index of highest log entry known to be committed
	LastApplied int // index of highest log entry applied to state machine

	readLease time.Time
}

func (s *raftState) getRole() pb.RaftRole {
	s.RLock()
	defer s.RUnlock()
	return s.Role
}

// maybe we could cache this field in Raft, to avoid lock
func (s *raftState) getTerm() int32 {
	s.RLock()
	defer s.RUnlock()
	return s.CurrentTerm
}

func (s *raftState) getCommited() int {
	s.RLock()
	defer s.RUnlock()
	return s.CommitIndex
}

const (
	currentTerm = "currentTerm"
	votedFor    = "votedFor"
	lastApplied = "lastApplied"
)

func (s *raftState) persist() {
	s.persister.StoreInt32(currentTerm, s.CurrentTerm)
	s.persister.StoreInt32(votedFor, s.VotedFor)
}

func (s *raftState) restore() {
	var ok bool
	s.CurrentTerm, ok = s.persister.LoadInt32(currentTerm)
	s.VotedFor, ok = s.persister.LoadInt32(votedFor)
	if !ok {
		s.VotedFor = -1
	}
	apply, ok := s.persister.LoadInt32(lastApplied)
	if ok {
		s.LastApplied = int(apply)
	}
}

func (s *raftState) commitUntil(index int) {
	s.CommitIndex = index
}

func (s *raftState) applyOne() {
	s.LastApplied++
	s.persister.StoreInt32(lastApplied, int32(s.LastApplied))
}

func (s *raftState) checkNewTerm(newTerm int32) bool {
	if s.getTerm() < newTerm {
		s.becomeFollower(-1, newTerm)
		return true
	}
	return false
}
func (s *raftState) checkFollowerCommit(leaderCommit int) bool {
	s.Lock()
	defer s.Unlock()
	// Have uncommited log, but leader has commited them
	if s.CommitIndex < leaderCommit && s.CommitIndex < s.log.LastIndex() {
		s.commitUntil(minInt(leaderCommit, int(s.log.LastIndex())))
		s.checkApply()
		return true
	}
	return false
}

// any command to apply
func (s *raftState) checkApply() {
	old := s.LastApplied
	log := s.log
	for s.CommitIndex > s.LastApplied {
		s.applyOne()
		if log.LastIndex() < s.LastApplied {
			panic(s)
		}
		entry := log.At(s.LastApplied)
		if entry == nil {
			break
		}
		msg := &ApplyMsg{Command: entry}
		s.applyCh <- msg
	}
	if s.LastApplied > old {
		glog.V(utils.VDebug).Infof("%s Applyed commands until index=%d", s.String(), s.LastApplied)
	}
}

func (s *raftState) changeRole(role pb.RaftRole) {
	s.Lock()
	defer s.Unlock()
	s.Role = role

}

func (s *raftState) becomeFollowerUnlocked(candidateID int32, term int32) {
	s.CurrentTerm = term
	s.VotedFor = candidateID
	s.Role = pb.RaftRole_Follower
	s.persist()
}

func (s *raftState) becomeFollower(candidateID int32, term int32) {
	s.Lock()
	defer s.Unlock()
	s.CurrentTerm = term
	s.VotedFor = candidateID
	s.Role = pb.RaftRole_Follower
	s.MatchIndex = nil
	s.NextIndex = nil
	s.persist()
}

func (s *raftState) becomeCandidate() int32 {
	s.Lock()
	defer s.Unlock()
	s.Role = pb.RaftRole_Candidate
	s.CurrentTerm++
	s.VotedFor = int32(s.Me)
	s.MatchIndex = nil
	s.NextIndex = nil
	s.persist()
	return s.CurrentTerm
}

func (s *raftState) becomeLeader() {
	s.Lock()
	defer s.Unlock()
	s.Role = pb.RaftRole_Leader
	s.MatchIndex = make([]int, s.NumPeers)
	s.NextIndex = make([]int, s.NumPeers)
	lastIndex := s.log.LastIndex()
	for i := range s.NextIndex {
		s.NextIndex[i] = lastIndex + 1
	}
}

func (s *raftState) replicatedToPeer(peer int, index int) {
	s.Lock()
	defer s.Unlock()
	s.MatchIndex[peer] = maxInt(index, s.MatchIndex[peer])
	s.NextIndex[peer] = s.MatchIndex[peer] + 1
	if s.checkLeaderCommit() {
		s.checkApply()
	}
}

// update  commit index, use matchIndex from peers
// NOTE: It's not thread-safe, should be synchronized by external lock
func (s *raftState) checkLeaderCommit() bool {
	matches := make([]int, 0, len(s.MatchIndex))
	for _, x := range s.MatchIndex {
		matches = append(matches, x)
	}
	sort.Sort(sort.Reverse(sort.IntSlice(matches)))
	newC := matches[s.majority()-1]
	if newC > s.CommitIndex {
		s.commitUntil(newC)
		glog.V(utils.VDebug).Infof("%s Leader update commitIndex: %d", s.String(), newC)
		return true
	}
	return false
}

func (s *raftState) retreatForPeer(peer int) int {
	s.Lock()
	defer s.Unlock()
	idx := &s.NextIndex[peer]
	log := s.log
	*idx = minInt(*idx, log.LastIndex())
	oldTerm := log.At(*idx).Term
	for *idx > 0 && log.At(*idx).Term == oldTerm {
		*idx--
	}
	*idx = maxInt(*idx, 1)
	return *idx
}

func (s *raftState) toReplicate(peer int) int {
	return s.NextIndex[peer]
}

// If lease is granted in this term, and later than the old one, extend the lease
func (s *raftState) updateReadLease(term int32, lease time.Time) {
	if s.readLease.Before(time.Now()) && term == s.CurrentTerm && lease.After(s.readLease) {
		s.readLease = lease
		glog.Infof("%s Update read lease to %v", s.String(), lease)
	}
}

func (s *raftState) String() string {
	return fmt.Sprintf("Raft<%d:%d>", s.Me, s.CurrentTerm)
}

func (s *raftState) dump(writer io.Writer) {
	state := map[string]interface{}{
		"meta":         s,
		"LogLastIndex": s.log.LastIndex(),
	}
	buf, err := json.MarshalIndent(state, "", "\t")
	if err != nil {
		panic(err)
	}
	writer.Write(buf)
}

func (s *raftState) majority() int {
	return s.NumPeers/2 + 1
}

func makeRaftState(log *store.LogStorage, persister store.Persister, applyCh chan<- *ApplyMsg, numPeers int, me int) *raftState {
	s := &raftState{
		log:         log,
		applyCh:     applyCh,
		NumPeers:    numPeers,
		Me:          me,
		CurrentTerm: 0,
		VotedFor:    -1,
		Role:        pb.RaftRole_Follower,
		readLease:   time.Now(),
		persister:   persister,
	}
	s.restore()
	glog.Infof("Restore raft state: CurrentTerm: %d, CommitIndex: %d, LastApplied: %d", s.CurrentTerm, s.CommitIndex, s.LastApplied)
	return s
}
<|endoftext|>


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

// To run these tests, set GCLOUD_TESTS_GOLANG_PROJECT_ID env var to your GCP projectID

package compute

import (
	"context"
	"fmt"
	"testing"

	"cloud.google.com/go/internal/testutil"
	"cloud.google.com/go/internal/uid"
	"google.golang.org/api/iterator"
	computepb "google.golang.org/genproto/googleapis/cloud/compute/v1"
	"google.golang.org/protobuf/proto"
)

var projectId = testutil.ProjID()
var defaultZone = "us-central1-a"

func TestCreateGetListInstance(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping smoke test in short mode")
	}
	space := uid.NewSpace("gogapic", nil)
	name := space.New()
	ctx := context.Background()
	c, err := NewInstancesRESTClient(ctx)
	if err != nil {
		t.Fatal(err)
	}
	zonesClient, err := NewZoneOperationsRESTClient(ctx)
	if err != nil {
		t.Fatal(err)
	}
	createRequest := &computepb.InsertInstanceRequest{
		Project: projectId,
		Zone:    defaultZone,
		InstanceResource: &computepb.Instance{
			Name:        &name,
			Description: proto.String("тест"),
			MachineType: proto.String(fmt.Sprintf("https://www.googleapis.com/compute/v1/projects/%s/zones/%s/machineTypes/n1-standard-1", projectId, defaultZone)),
			Disks: []*computepb.AttachedDisk{
				{
					AutoDelete: proto.Bool(true),
					Boot:       proto.Bool(true),
					Type:       computepb.AttachedDisk_PERSISTENT.Enum(),
					InitializeParams: &computepb.AttachedDiskInitializeParams{
						SourceImage: proto.String("projects/debian-cloud/global/images/family/debian-10"),
					},
				},
			},
			NetworkInterfaces: []*computepb.NetworkInterface{
				{
					AccessConfigs: []*computepb.AccessConfig{
						{
							Name: proto.String("default"),
						},
					},
				},
			},
		},
	}

	insert, err := c.Insert(ctx, createRequest)
	if err != nil {
		t.Fatal(err)
	}

	waitZonalRequest := &computepb.WaitZoneOperationRequest{
		Project:   projectId,
		Zone:      defaultZone,
		Operation: insert.Proto().GetName(),
	}
	_, err = zonesClient.Wait(ctx, waitZonalRequest)
	if err != nil {
		t.Error(err)
	}
	defer ForceDeleteInstance(ctx, name, c)

	getRequest := &computepb.GetInstanceRequest{
		Project:  projectId,
		Zone:     defaultZone,
		Instance: name,
	}
	get, err := c.Get(ctx, getRequest)
	if err != nil {
		t.Error(err)
	}
	if get.GetName() != name {
		t.Fatal(fmt.Sprintf("expected instance name: %s, got: %s", name, get.GetName()))
	}
	if get.GetDescription() != "тест" {
		t.Fatal(fmt.Sprintf("expected instance description: %s, got: %s", "тест", get.GetDescription()))
	}
	listRequest := &computepb.ListInstancesRequest{
		Project: projectId,
		Zone:    defaultZone,
	}

	itr := c.List(ctx, listRequest)
	if err != nil {
		t.Error(err)
	}
	found := false
	element, err := itr.Next()
	for err == nil {
		if element.GetName() == name {
			found = true
		}
		element, err = itr.Next()
	}
	if err != nil && err != iterator.Done {
		t.Fatal(err)
	}
	if !found {
		t.Error("Couldn't find the instance in list response")
	}

	deleteInstanceRequest := &computepb.DeleteInstanceRequest{
		Project:  projectId,
		Zone:     defaultZone,
		Instance: name,
	}
	_, err = c.Delete(ctx, deleteInstanceRequest)
	if err != nil {
		t.Error(err)
	}
}

func ForceDeleteInstance(ctx context.Context, name string, client *InstancesClient) {
	deleteInstanceRequest := &computepb.DeleteInstanceRequest{
		Project:  projectId,
		Zone:     defaultZone,
		Instance: name,
	}
	client.Delete(ctx, deleteInstanceRequest)
}

func TestCreateGetRemoveSecurityPolicies(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping smoke test in short mode")
	}
	space := uid.NewSpace("gogapic", nil)
	name := space.New()
	ctx := context.Background()
	c, err := NewSecurityPoliciesRESTClient(ctx)
	if err != nil {
		t.Fatal(err)
	}
	globalCLient, err := NewGlobalOperationsRESTClient(ctx)
	if err != nil {
		t.Fatal(err)
	}
	action := "allow"
	matcher := &computepb.SecurityPolicyRuleMatcher{
		Config: &computepb.SecurityPolicyRuleMatcherConfig{
			SrcIpRanges: []string{
				"*",
			},
		},
		VersionedExpr: computepb.SecurityPolicyRuleMatcher_SRC_IPS_V1.Enum(),
	}
	securityPolicyRule := &computepb.SecurityPolicyRule{
		Action:      &action,
		Priority:    proto.Int32(0),
		Description: proto.String("test rule"),
		Match:       matcher,
	}
	securityPolicyRuleDefault := &computepb.SecurityPolicyRule{
		Action:      &action,
		Priority:    proto.Int32(2147483647),
		Description: proto.String("default rule"),
		Match:       matcher,
	}
	insertRequest := &computepb.InsertSecurityPolicyRequest{
		Project: projectId,
		SecurityPolicyResource: &computepb.SecurityPolicy{
			Name: &name,
			Rules: []*computepb.SecurityPolicyRule{
				securityPolicyRule,
				securityPolicyRuleDefault,
			},
		},
	}
	insert, err := c.Insert(ctx, insertRequest)
	if err != nil {
		t.Fatal(err)
	}

	waitGlobalRequest := &computepb.WaitGlobalOperationRequest{
		Project:   projectId,
		Operation: insert.Proto().GetName(),
	}
	_, err = globalCLient.Wait(ctx, waitGlobalRequest)
	if err != nil {
		t.Error(err)
	}
	defer ForceDeleteSecurityPolicy(ctx, name, c)

	removeRuleRequest := &computepb.RemoveRuleSecurityPolicyRequest{
		Priority:       proto.Int32(0),
		Project:        projectId,
		SecurityPolicy: name,
	}

	rule, err := c.RemoveRule(ctx, removeRuleRequest)
	if err != nil {
		t.Error(err)
	}
	waitGlobalRequestRemove := &computepb.WaitGlobalOperationRequest{
		Project:   projectId,
		Operation: rule.Proto().GetName(),
	}
	_, err = globalCLient.Wait(ctx, waitGlobalRequestRemove)
	if err != nil {
		t.Error(err)
	}

	getRequest := &computepb.GetSecurityPolicyRequest{
		Project:        projectId,
		SecurityPolicy: name,
	}
	get, err := c.Get(ctx, getRequest)
	if err != nil {
		t.Error(err)
	}
	if len(get.GetRules()) != 1 {
		t.Fatal(fmt.Sprintf("expected count for rules: %d, got: %d", 1, len(get.GetRules())))
	}

	deleteRequest := &computepb.DeleteSecurityPolicyRequest{
		Project:        projectId,
		SecurityPolicy: name,
	}
	_, err = c.Delete(ctx, deleteRequest)
	if err != nil {
		t.Error(err)
	}
}

func ForceDeleteSecurityPolicy(ctx context.Context, name string, client *SecurityPoliciesClient) {
	deleteRequest := &computepb.DeleteSecurityPolicyRequest{
		Project:        projectId,
		SecurityPolicy: name,
	}
	client.Delete(ctx, deleteRequest)
}

func TestPaginationWithMaxRes(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping smoke test in short mode")
	}
	ctx := context.Background()
	c, err := NewAcceleratorTypesRESTClient(ctx)
	if err != nil {
		t.Fatal(err)
	}
	req := &computepb.ListAcceleratorTypesRequest{
		Project:    projectId,
		Zone:       defaultZone,
		MaxResults: proto.Uint32(1),
	}
	itr := c.List(ctx, req)

	found := false
	element, err := itr.Next()
	for err == nil {
		if element.GetName() == "nvidia-tesla-t4" {
			found = true
		}
		element, err = itr.Next()
	}
	if err != iterator.Done {
		t.Fatal(err)
	}
	if !found {
		t.Error("Couldn't find the accelerator in the response")
	}
}

func TestPaginationDefault(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping smoke test in short mode")
	}
	ctx := context.Background()
	c, err := NewAcceleratorTypesRESTClient(ctx)
	if err != nil {
		t.Fatal(err)
	}
	req := &computepb.ListAcceleratorTypesRequest{
		Project: projectId,
		Zone:    defaultZone,
	}
	itr := c.List(ctx, req)

	found := false
	element, err := itr.Next()
	for err == nil {
		if element.GetName() == "nvidia-tesla-t4" {
			found = true
			break
		}
		element, err = itr.Next()
	}
	if err != iterator.Done {
		t.Fatal(err)
	}
	if !found {
		t.Error("Couldn't find the accelerator in the response")
	}
}

func TestPaginationMapResponse(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping smoke test in short mode")
	}
	ctx := context.Background()
	c, err := NewAcceleratorTypesRESTClient(ctx)
	if err != nil {
		t.Fatal(err)
	}
	req := &computepb.AggregatedListAcceleratorTypesRequest{
		Project: projectId,
	}
	itr := c.AggregatedList(ctx, req)

	found := false
	element, err := itr.Next()
	for err == nil {
		if element.Key == "zones/us-central1-a" {
			types := element.Value.GetAcceleratorTypes()
			for _, item := range types {
				if item.GetName() == "nvidia-tesla-t4" {
					found = true
					break
				}
			}
		}
		element, err = itr.Next()
	}
	if err != iterator.Done {
		t.Fatal(err)
	}
	if !found {
		t.Error("Couldn't find the accelerator in the response")
	}
}

func TestPaginationMapResponseMaxRes(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping smoke test in short mode")
	}
	ctx := context.Background()
	c, err := NewAcceleratorTypesRESTClient(ctx)
	if err != nil {
		t.Fatal(err)
	}
	req := &computepb.AggregatedListAcceleratorTypesRequest{
		Project:    projectId,
		MaxResults: proto.Uint32(10),
	}
	itr := c.AggregatedList(ctx, req)
	found := false
	element, err := itr.Next()
	for err == nil {
		if element.Key == "zones/us-central1-a" {
			types := element.Value.GetAcceleratorTypes()
			for _, item := range types {
				if item.GetName() == "nvidia-tesla-t4" {
					found = true
					break
				}
			}
		}
		element, err = itr.Next()
	}
	if err != iterator.Done {
		t.Fatal(err)
	}
	if !found {
		t.Error("Couldn't find the accelerator in the response")
	}
}
<commit_msg>tests(compute): fix failing test (#4549)<commit_after>// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// To run these tests, set GCLOUD_TESTS_GOLANG_PROJECT_ID env var to your GCP projectID

package compute

import (
	"context"
	"fmt"
	"testing"

	"cloud.google.com/go/internal/testutil"
	"cloud.google.com/go/internal/uid"
	"google.golang.org/api/iterator"
	computepb "google.golang.org/genproto/googleapis/cloud/compute/v1"
	"google.golang.org/protobuf/proto"
)

var projectId = testutil.ProjID()
var defaultZone = "us-central1-a"

func TestCreateGetListInstance(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping smoke test in short mode")
	}
	space := uid.NewSpace("gogapic", nil)
	name := space.New()
	ctx := context.Background()
	c, err := NewInstancesRESTClient(ctx)
	if err != nil {
		t.Fatal(err)
	}
	zonesClient, err := NewZoneOperationsRESTClient(ctx)
	if err != nil {
		t.Fatal(err)
	}
	createRequest := &computepb.InsertInstanceRequest{
		Project: projectId,
		Zone:    defaultZone,
		InstanceResource: &computepb.Instance{
			Name:        &name,
			Description: proto.String("тест"),
			MachineType: proto.String(fmt.Sprintf("https://www.googleapis.com/compute/v1/projects/%s/zones/%s/machineTypes/n1-standard-1", projectId, defaultZone)),
			Disks: []*computepb.AttachedDisk{
				{
					AutoDelete: proto.Bool(true),
					Boot:       proto.Bool(true),
					Type:       computepb.AttachedDisk_PERSISTENT.Enum(),
					InitializeParams: &computepb.AttachedDiskInitializeParams{
						SourceImage: proto.String("projects/debian-cloud/global/images/family/debian-10"),
					},
				},
			},
			NetworkInterfaces: []*computepb.NetworkInterface{
				{
					AccessConfigs: []*computepb.AccessConfig{
						{
							Name: proto.String("default"),
						},
					},
				},
			},
		},
	}

	insert, err := c.Insert(ctx, createRequest)
	if err != nil {
		t.Fatal(err)
	}

	waitZonalRequest := &computepb.WaitZoneOperationRequest{
		Project:   projectId,
		Zone:      defaultZone,
		Operation: insert.Proto().GetName(),
	}
	_, err = zonesClient.Wait(ctx, waitZonalRequest)
	if err != nil {
		t.Error(err)
	}
	defer ForceDeleteInstance(ctx, name, c)

	getRequest := &computepb.GetInstanceRequest{
		Project:  projectId,
		Zone:     defaultZone,
		Instance: name,
	}
	get, err := c.Get(ctx, getRequest)
	if err != nil {
		t.Error(err)
	}
	if get.GetName() != name {
		t.Fatal(fmt.Sprintf("expected instance name: %s, got: %s", name, get.GetName()))
	}
	if get.GetDescription() != "тест" {
		t.Fatal(fmt.Sprintf("expected instance description: %s, got: %s", "тест", get.GetDescription()))
	}
	listRequest := &computepb.ListInstancesRequest{
		Project: projectId,
		Zone:    defaultZone,
	}

	itr := c.List(ctx, listRequest)
	if err != nil {
		t.Error(err)
	}
	found := false
	element, err := itr.Next()
	for err == nil {
		if element.GetName() == name {
			found = true
		}
		element, err = itr.Next()
	}
	if err != nil && err != iterator.Done {
		t.Fatal(err)
	}
	if !found {
		t.Error("Couldn't find the instance in list response")
	}

	deleteInstanceRequest := &computepb.DeleteInstanceRequest{
		Project:  projectId,
		Zone:     defaultZone,
		Instance: name,
	}
	_, err = c.Delete(ctx, deleteInstanceRequest)
	if err != nil {
		t.Error(err)
	}
}

func ForceDeleteInstance(ctx context.Context, name string, client *InstancesClient) {
	deleteInstanceRequest := &computepb.DeleteInstanceRequest{
		Project:  projectId,
		Zone:     defaultZone,
		Instance: name,
	}
	client.Delete(ctx, deleteInstanceRequest)
}

func TestCreateGetRemoveSecurityPolicies(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping smoke test in short mode")
	}
	space := uid.NewSpace("gogapic", nil)
	name := space.New()
	ctx := context.Background()
	c, err := NewSecurityPoliciesRESTClient(ctx)
	if err != nil {
		t.Fatal(err)
	}
	globalCLient, err := NewGlobalOperationsRESTClient(ctx)
	if err != nil {
		t.Fatal(err)
	}
	action := "allow"
	matcher := &computepb.SecurityPolicyRuleMatcher{
		Config: &computepb.SecurityPolicyRuleMatcherConfig{
			SrcIpRanges: []string{
				"*",
			},
		},
		VersionedExpr: computepb.SecurityPolicyRuleMatcher_SRC_IPS_V1.Enum(),
	}
	securityPolicyRule := &computepb.SecurityPolicyRule{
		Action:      &action,
		Priority:    proto.Int32(0),
		Description: proto.String("test rule"),
		Match:       matcher,
	}
	securityPolicyRuleDefault := &computepb.SecurityPolicyRule{
		Action:      &action,
		Priority:    proto.Int32(2147483647),
		Description: proto.String("default rule"),
		Match:       matcher,
	}
	insertRequest := &computepb.InsertSecurityPolicyRequest{
		Project: projectId,
		SecurityPolicyResource: &computepb.SecurityPolicy{
			Name: &name,
			Rules: []*computepb.SecurityPolicyRule{
				securityPolicyRule,
				securityPolicyRuleDefault,
			},
		},
	}
	insert, err := c.Insert(ctx, insertRequest)
	if err != nil {
		t.Fatal(err)
	}

	waitGlobalRequest := &computepb.WaitGlobalOperationRequest{
		Project:   projectId,
		Operation: insert.Proto().GetName(),
	}
	_, err = globalCLient.Wait(ctx, waitGlobalRequest)
	if err != nil {
		t.Error(err)
	}
	defer ForceDeleteSecurityPolicy(ctx, name, c)

	removeRuleRequest := &computepb.RemoveRuleSecurityPolicyRequest{
		Priority:       proto.Int32(0),
		Project:        projectId,
		SecurityPolicy: name,
	}

	rule, err := c.RemoveRule(ctx, removeRuleRequest)
	if err != nil {
		t.Error(err)
	}
	waitGlobalRequestRemove := &computepb.WaitGlobalOperationRequest{
		Project:   projectId,
		Operation: rule.Proto().GetName(),
	}
	_, err = globalCLient.Wait(ctx, waitGlobalRequestRemove)
	if err != nil {
		t.Error(err)
	}

	getRequest := &computepb.GetSecurityPolicyRequest{
		Project:        projectId,
		SecurityPolicy: name,
	}
	get, err := c.Get(ctx, getRequest)
	if err != nil {
		t.Error(err)
	}
	if len(get.GetRules()) != 1 {
		t.Fatal(fmt.Sprintf("expected count for rules: %d, got: %d", 1, len(get.GetRules())))
	}

	deleteRequest := &computepb.DeleteSecurityPolicyRequest{
		Project:        projectId,
		SecurityPolicy: name,
	}
	_, err = c.Delete(ctx, deleteRequest)
	if err != nil {
		t.Error(err)
	}
}

func ForceDeleteSecurityPolicy(ctx context.Context, name string, client *SecurityPoliciesClient) {
	deleteRequest := &computepb.DeleteSecurityPolicyRequest{
		Project:        projectId,
		SecurityPolicy: name,
	}
	client.Delete(ctx, deleteRequest)
}

func TestPaginationWithMaxRes(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping smoke test in short mode")
	}
	ctx := context.Background()
	c, err := NewAcceleratorTypesRESTClient(ctx)
	if err != nil {
		t.Fatal(err)
	}
	req := &computepb.ListAcceleratorTypesRequest{
		Project:    projectId,
		Zone:       defaultZone,
		MaxResults: proto.Uint32(1),
	}
	itr := c.List(ctx, req)

	found := false
	element, err := itr.Next()
	for err == nil {
		if element.GetName() == "nvidia-tesla-t4" {
			found = true
			break
		}
		element, err = itr.Next()
	}
	if err != nil && err != iterator.Done {
		t.Fatal(err)
	}
	if !found {
		t.Error("Couldn't find the accelerator in the response")
	}
}

func TestPaginationDefault(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping smoke test in short mode")
	}
	ctx := context.Background()
	c, err := NewAcceleratorTypesRESTClient(ctx)
	if err != nil {
		t.Fatal(err)
	}
	req := &computepb.ListAcceleratorTypesRequest{
		Project: projectId,
		Zone:    defaultZone,
	}
	itr := c.List(ctx, req)

	found := false
	element, err := itr.Next()
	for err == nil {
		if element.GetName() == "nvidia-tesla-t4" {
			found = true
			break
		}
		element, err = itr.Next()
	}
	if err != nil && err != iterator.Done {
		t.Fatal(err)
	}
	if !found {
		t.Error("Couldn't find the accelerator in the response")
	}
}

func TestPaginationMapResponse(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping smoke test in short mode")
	}
	ctx := context.Background()
	c, err := NewAcceleratorTypesRESTClient(ctx)
	if err != nil {
		t.Fatal(err)
	}
	req := &computepb.AggregatedListAcceleratorTypesRequest{
		Project: projectId,
	}
	itr := c.AggregatedList(ctx, req)

	found := false
	element, err := itr.Next()
	for err == nil {
		if element.Key == "zones/us-central1-a" {
			types := element.Value.GetAcceleratorTypes()
			for _, item := range types {
				if item.GetName() == "nvidia-tesla-t4" {
					found = true
					break
				}
			}
		}
		element, err = itr.Next()
	}
	if err != iterator.Done {
		t.Fatal(err)
	}
	if !found {
		t.Error("Couldn't find the accelerator in the response")
	}
}

func TestPaginationMapResponseMaxRes(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping smoke test in short mode")
	}
	ctx := context.Background()
	c, err := NewAcceleratorTypesRESTClient(ctx)
	if err != nil {
		t.Fatal(err)
	}
	req := &computepb.AggregatedListAcceleratorTypesRequest{
		Project:    projectId,
		MaxResults: proto.Uint32(10),
	}
	itr := c.AggregatedList(ctx, req)
	found := false
	element, err := itr.Next()
	for err == nil {
		if element.Key == "zones/us-central1-a" {
			types := element.Value.GetAcceleratorTypes()
			for _, item := range types {
				if item.GetName() == "nvidia-tesla-t4" {
					found = true
					break
				}
			}
		}
		element, err = itr.Next()
	}
	if err != iterator.Done {
		t.Fatal(err)
	}
	if !found {
		t.Error("Couldn't find the accelerator in the response")
	}
}
<|endoftext|>


<commit_before>// The term package implements rendering HTML content for the terminal.
package term

import (
	"bytes"
	"errors"
	"io"
	"runtime"
	"unicode"

	"github.com/PuerkitoBio/goquery"
	"golang.org/x/net/html"
)

// Renderer implements a rendering HTML content in a stripped down format
// suitable for the terminal.
type Renderer struct {
	// Width sets the assumed width of the terminal, which aids in wrapping
	// text. If == 0, then the width is treated as 80. If < 0, then the width is
	// treated as unbounded.
	Width int
	// TabSize indicates the number of characters that the terminal uses to
	// render a tab. If <= 0, then tabs are assumed to have a size of 8.
	TabSize int
}

var (
	stop         = errors.New("stop")          // Stop walking without erroring.
	skipChildren = errors.New("skip children") // Skip all child nodes.
	skipText     = errors.New("skip text")     // Skip child text nodes.
)

type walker func(node *html.Node, entering bool) error

func walk(node *html.Node, cb walker) error {
	err := cb(node, true)
	if err != nil && err != skipChildren && err != skipText {
		return err
	}
	if err != skipChildren {
		for c := node.FirstChild; c != nil; c = c.NextSibling {
			if c.Type == html.TextNode && err == skipText {
				continue
			}
			if err := walk(c, cb); err != nil {
				return err
			}
		}
	}
	if err = cb(node, false); err != nil && err != skipChildren && err != skipText {
		return err
	}
	return err
}

// Facilitates wrapped text.
type writer struct {
	bytes.Buffer
	w io.Writer

	// If wrapping, wrap to this width.
	width int

	// Size of tab character.
	tabSize int
}

func (w *writer) Flush() error {
	b := w.Buffer.Bytes()
	defer w.Buffer.Reset()

	b = bytes.TrimSpace(b)
	var paragraphs [][]rune
	for i := 0; i < len(b); {
		if b[i] == '\n' {
			// Bounds check unneeded; if b[i], being a \n, was at the end, it
			// would have been trimmed.
			if b[i+1] == '\n' {
				// New paragraph.
				paragraphs = append(paragraphs, []rune(string(b[:i])))
				i += 2
				for b[i] == '\n' {
					// Collapse extra newlines.
					i++
				}
				b = b[i:]
				i = 0
				continue
			} else {
				// Unwrap.
				b[i] = ' '
			}
		}
		i++
	}
	if w.width > 0 {
		for j, p := range paragraphs {
			for i := 0; i < len(p); {
				n := i + w.width
				if n+1 >= len(p) {
					break
				}
				for n > i && !unicode.IsSpace(p[n]) {
					n--
				}
				if n <= i {
					// Long word.
					n = i + w.width
					p = append(p, 0)
					copy(p[n+1:], p[n:])
					p[n] = '\n'
					i = n + 1
				} else {
					p[n] = '\n'
					i = n + 1
				}
			}
			paragraphs[j] = p
		}
		if runtime.GOOS == "windows" {
			// Remove newlines that are inserted just after the edge of the
			// width. This prevents the terminal from incorrectly producing an
			// extra gap when wrapping. This removes the newline entirely, so it
			// is assumed that the wrapping will separate the previous word from
			// the next word.
			for i, p := range paragraphs {
				for h, i := 0, 0; i < len(p); i++ {
					if p[i] != '\n' {
						continue
					}
					if i-h >= w.width {
						copy(p[i:], p[i+1:])
						p = p[:len(p)-1]
						h = i
						i--
					} else {
						h = i + 1
					}
				}
				paragraphs[i] = p
			}
		}
	}
	for i, p := range paragraphs {
		if i > 0 {
			w.w.Write([]byte{'\n', '\n'})
		}
		if _, err := w.w.Write([]byte(string(p))); err != nil {
			return err
		}
	}
	return nil
}

func (r Renderer) Render(w io.Writer, s *goquery.Selection) error {
	buf := &writer{
		w:       w,
		width:   r.Width,
		tabSize: r.TabSize,
	}
	for _, node := range s.Nodes {
		if node.Type == html.TextNode {
			continue
		}
		err := walk(node, func(node *html.Node, entering bool) error {
			switch node.Type {
			case html.ErrorNode:
			case html.TextNode:
				if entering {
					buf.WriteString(node.Data)
				}
			case html.DocumentNode:
			case html.ElementNode:
				h := handlers[elementMatcher{node.Data, entering}]
				if h != nil {
					return h(buf, node)
				}
			case html.CommentNode:
			case html.DoctypeNode:
			case html.RawNode:
			}
			return nil
		})
		if err != nil && err != stop {
			return err
		}
	}
	return buf.Flush()
}

type elementMatcher struct {
	data     string
	entering bool
}

type nodeHandlers map[elementMatcher]func(w *writer, node *html.Node) error

func isElement(node *html.Node, tag string) bool {
	return node.Type == html.ElementNode && node.Data == tag
}

var handlers = nodeHandlers{
	{"code", true}: func(w *writer, node *html.Node) error {
		if isElement(node.Parent, "pre") { // May have syntax: `class="language-*"`
			// Block
			w.WriteString(node.FirstChild.Data)
			return skipChildren
		} else {
			// Inline
			return w.WriteByte('`')
		}
	},
	{"code", false}: func(w *writer, node *html.Node) error {
		if isElement(node.Parent, "pre") {
			// Block
		} else {
			// Inline
			return w.WriteByte('`')
		}
		return nil
	},
	{"section", true}: func(w *writer, node *html.Node) error {
		return skipText
	},
	{"p", true}: func(w *writer, node *html.Node) error {
		_, err := w.WriteString("\n\n")
		return err
	},
	{"p", false}: func(w *writer, node *html.Node) error {
		_, err := w.WriteString("\n\n")
		return err
	},
}
<commit_msg>Fix missing paragraph.<commit_after>// The term package implements rendering HTML content for the terminal.
package term

import (
	"bytes"
	"errors"
	"io"
	"runtime"
	"unicode"

	"github.com/PuerkitoBio/goquery"
	"golang.org/x/net/html"
)

// Renderer implements a rendering HTML content in a stripped down format
// suitable for the terminal.
type Renderer struct {
	// Width sets the assumed width of the terminal, which aids in wrapping
	// text. If == 0, then the width is treated as 80. If < 0, then the width is
	// treated as unbounded.
	Width int
	// TabSize indicates the number of characters that the terminal uses to
	// render a tab. If <= 0, then tabs are assumed to have a size of 8.
	TabSize int
}

var (
	stop         = errors.New("stop")          // Stop walking without erroring.
	skipChildren = errors.New("skip children") // Skip all child nodes.
	skipText     = errors.New("skip text")     // Skip child text nodes.
)

type walker func(node *html.Node, entering bool) error

func walk(node *html.Node, cb walker) error {
	err := cb(node, true)
	if err != nil && err != skipChildren && err != skipText {
		return err
	}
	if err != skipChildren {
		for c := node.FirstChild; c != nil; c = c.NextSibling {
			if c.Type == html.TextNode && err == skipText {
				continue
			}
			if err := walk(c, cb); err != nil {
				return err
			}
		}
	}
	if err = cb(node, false); err != nil && err != skipChildren && err != skipText {
		return err
	}
	return err
}

// Facilitates wrapped text.
type writer struct {
	bytes.Buffer
	w io.Writer

	// If wrapping, wrap to this width.
	width int

	// Size of tab character.
	tabSize int
}

func (w *writer) Flush() error {
	b := w.Buffer.Bytes()
	defer w.Buffer.Reset()

	b = bytes.TrimSpace(b)
	var paragraphs [][]rune
	for i := 0; i < len(b); {
		if b[i] == '\n' {
			// Bounds check unneeded; if b[i], being a \n, was at the end, it
			// would have been trimmed.
			if b[i+1] == '\n' {
				// New paragraph.
				paragraphs = append(paragraphs, []rune(string(b[:i])))
				i += 2
				for b[i] == '\n' {
					// Collapse extra newlines.
					i++
				}
				b = b[i:]
				i = 0
				continue
			} else {
				// Unwrap.
				b[i] = ' '
			}
		}
		i++
	}
	paragraphs = append(paragraphs, []rune(string(b)))
	if w.width > 0 {
		for j, p := range paragraphs {
			for i := 0; i < len(p); {
				n := i + w.width
				if n+1 >= len(p) {
					break
				}
				for n > i && !unicode.IsSpace(p[n]) {
					n--
				}
				if n <= i {
					// Long word.
					n = i + w.width
					p = append(p, 0)
					copy(p[n+1:], p[n:])
					p[n] = '\n'
					i = n + 1
				} else {
					p[n] = '\n'
					i = n + 1
				}
			}
			paragraphs[j] = p
		}
		if runtime.GOOS == "windows" {
			// Remove newlines that are inserted just after the edge of the
			// width. This prevents the terminal from incorrectly producing an
			// extra gap when wrapping. This removes the newline entirely, so it
			// is assumed that the wrapping will separate the previous word from
			// the next word.
			for i, p := range paragraphs {
				for h, i := 0, 0; i < len(p); i++ {
					if p[i] != '\n' {
						continue
					}
					if i-h >= w.width {
						copy(p[i:], p[i+1:])
						p = p[:len(p)-1]
						h = i
						i--
					} else {
						h = i + 1
					}
				}
				paragraphs[i] = p
			}
		}
	}
	for i, p := range paragraphs {
		if i > 0 {
			w.w.Write([]byte{'\n', '\n'})
		}
		if _, err := w.w.Write([]byte(string(p))); err != nil {
			return err
		}
	}
	return nil
}

func (r Renderer) Render(w io.Writer, s *goquery.Selection) error {
	buf := &writer{
		w:       w,
		width:   r.Width,
		tabSize: r.TabSize,
	}
	for _, node := range s.Nodes {
		if node.Type == html.TextNode {
			continue
		}
		err := walk(node, func(node *html.Node, entering bool) error {
			switch node.Type {
			case html.ErrorNode:
			case html.TextNode:
				if entering {
					buf.WriteString(node.Data)
				}
			case html.DocumentNode:
			case html.ElementNode:
				h := handlers[elementMatcher{node.Data, entering}]
				if h != nil {
					return h(buf, node)
				}
			case html.CommentNode:
			case html.DoctypeNode:
			case html.RawNode:
			}
			return nil
		})
		if err != nil && err != stop {
			return err
		}
	}
	return buf.Flush()
}

type elementMatcher struct {
	data     string
	entering bool
}

type nodeHandlers map[elementMatcher]func(w *writer, node *html.Node) error

func isElement(node *html.Node, tag string) bool {
	return node.Type == html.ElementNode && node.Data == tag
}

var handlers = nodeHandlers{
	{"code", true}: func(w *writer, node *html.Node) error {
		if isElement(node.Parent, "pre") { // May have syntax: `class="language-*"`
			// Block
			w.WriteString(node.FirstChild.Data)
			return skipChildren
		} else {
			// Inline
			return w.WriteByte('`')
		}
	},
	{"code", false}: func(w *writer, node *html.Node) error {
		if isElement(node.Parent, "pre") {
			// Block
		} else {
			// Inline
			return w.WriteByte('`')
		}
		return nil
	},
	{"section", true}: func(w *writer, node *html.Node) error {
		return skipText
	},
	{"p", true}: func(w *writer, node *html.Node) error {
		_, err := w.WriteString("\n\n")
		return err
	},
	{"p", false}: func(w *writer, node *html.Node) error {
		_, err := w.WriteString("\n\n")
		return err
	},
}
<|endoftext|>


<commit_before>// Copyright 2015 Yahoo Inc.
// Licensed under the terms of the Apache version 2.0 license. See LICENSE file for terms.

package main

import (
	"bufio"
	"github.com/ardielle/ardielle-go/rdl"
	"strings"
	"text/template"
)

type javaClientGenerator struct {
	registry rdl.TypeRegistry
	schema   *rdl.Schema
	name     string
	writer   *bufio.Writer
	err      error
	banner   string
	ns       string
	base     string
}

// GenerateJavaClient generates the client code to talk to the server
func GenerateJavaClient(banner string, schema *rdl.Schema, outdir string, ns string, base string, options []string) error {
	reg := rdl.NewTypeRegistry(schema)
	packageDir, err := javaGenerationDir(outdir, schema, ns)
	if err != nil {
		return err
	}

	cName := javaGenerationStringOptionSet(options, "clientclass")
	if cName == "" {
		cName = capitalize(string(schema.Name))
	}

	out, file, _, err := outputWriter(packageDir, cName, "Client.java")
	if err != nil {
		return err
	}
	gen := &javaClientGenerator{reg, schema, cName, out, nil, banner, ns, base}
	gen.processTemplate(javaClientTemplate)
	out.Flush()
	file.Close()
	if gen.err != nil {
		return gen.err
	}

	//ResourceException - the throawable wrapper for alternate return types
	out, file, _, err = outputWriter(packageDir, "ResourceException", ".java")
	if err != nil {
		return err
	}
	err = javaGenerateResourceException(schema, out, ns)
	out.Flush()
	file.Close()
	if err != nil {
		return err
	}

	//ResourceError - the default data object for an error
	out, file, _, err = outputWriter(packageDir, "ResourceError", ".java")
	if err != nil {
		return err
	}
	err = javaGenerateResourceError(schema, out, ns)
	out.Flush()
	file.Close()
	return err
}

func (gen *javaClientGenerator) processTemplate(templateSource string) error {
	commentFun := func(s string) string {
		return formatComment(s, 0, 80)
	}
	funcMap := template.FuncMap{
		"header":     func() string { return javaGenerationHeader(gen.banner) },
		"package":    func() string { return javaGenerationPackage(gen.schema, gen.ns) },
		"comment":    commentFun,
		"methodSig":  func(r *rdl.Resource) string { return gen.clientMethodSignature(r) },
		"methodBody": func(r *rdl.Resource) string { return gen.clientMethodBody(r) },
		"name":       func() string { return gen.name },
		"cName":      func() string { return capitalize(gen.name) },
		"lName":      func() string { return uncapitalize(gen.name) },
	}
	t := template.Must(template.New(gen.name).Funcs(funcMap).Parse(templateSource))
	return t.Execute(gen.writer, gen.schema)
}

func (gen *javaClientGenerator) resourcePath(r *rdl.Resource) string {
	path := r.Path
	i := strings.Index(path, "?")
	if i >= 0 {
		path = path[0:i]
	}
	return path
}

const javaClientTemplate = `{{header}}
package {{package}};
import com.yahoo.rdl.*;
import javax.ws.rs.client.*;
import javax.ws.rs.*;
import javax.ws.rs.core.*;

public class {{cName}}Client {
    WebTarget base;
    String credsHeader;
    String credsToken;

    public {{cName}}Client(String url) {
        base = ClientBuilder.newClient().target(url);
    }

    public {{cName}}Client addCredentials(String header, String token) {
        credsHeader = header;
        credsToken = token;
        return this;
    }
{{range .Resources}}
    {{methodSig .}} {
        {{methodBody .}}
    }
{{end}}
}
`

func (gen *javaClientGenerator) clientMethodSignature(r *rdl.Resource) string {
	reg := gen.registry
	returnType := javaType(reg, r.Type, false, "", "")
	methName, params := javaMethodName(reg, r)
	sparams := ""
	if len(params) > 0 {
		sparams = strings.Join(params, ", ")
	}
	return "public " + returnType + " " + methName + "(" + sparams + ")"
}

func (gen *javaClientGenerator) clientMethodBody(r *rdl.Resource) string {
	reg := gen.registry
	returnType := javaType(reg, r.Type, false, "", "")
	path := r.Path
	s := "WebTarget target = base.path(\"" + path + "\")"
	entityName := ""
	q := ""
	h := ""
	for _, in := range r.Inputs {
		iname := javaName(in.Name)
		if in.PathParam {
			s += "\n            .resolveTemplate(\"" + iname + "\", " + iname + ")"
		} else if in.QueryParam != "" {
			q += "\n        if (" + iname + " != null) {"
			q += "\n            target = target.queryParam(\"" + in.QueryParam + "\", " + iname + ");"
			q += "\n        }"
		} else if in.Header != "" {
			h += "\n        if (" + iname + " != null) {"
			h += "\n            invocationBuilder = invocationBuilder.header(\"" + in.Header + "\", " + iname + ");"
			h += "\n        }"
		} else { //the entity
			entityName = iname
		}
	}
	s += ";"
	if q != "" {
		s += q
	}
	s += "\n        Invocation.Builder invocationBuilder = target.request(\"application/json\");"
	if h != "" {
		s += h
	}
	s += "\n"
	switch r.Method {
	case "PUT", "POST":
		s += "        Response response = invocationBuilder." + strings.ToLower(r.Method) + "(Entity.entity(" + entityName + ", \"application/json\"));\n"
	default:
		s += "        Response response = invocationBuilder." + strings.ToLower(r.Method) + "();\n"
	}
	s += "        int code = response.getStatus();\n"
	s += "        String status = Response.Status.fromStatusCode(code).toString();\n"
	expected := "status.equals(\"" + r.Expected + "\")"
	for _, alt := range r.Alternatives {
		expected += "|| status.equals(\"" + alt + "\")"
	}
	s += "        if (" + expected + ") {\n"
	s += "            return response.readEntity(" + returnType + ".class);\n"
	if r.Exceptions != nil {
		for k, v := range r.Exceptions {
			s += "        } else if (status.equals(\"" + k + "\")) {\n"
			s += "            throw new ResourceException(code, response.readEntity(" + v.Type + ".class));\n"
		}
	}
	s += "        } else {\n"
	s += "            throw new ResourceException(code, response.readEntity(Object.class));\n"
	s += "        }\n"
	return s
}
<commit_msg>enhancements for java client generation<commit_after>// Copyright 2015 Yahoo Inc.
// Licensed under the terms of the Apache version 2.0 license. See LICENSE file for terms.

package main

import (
	"bufio"
	"github.com/ardielle/ardielle-go/rdl"
	"strings"
	"text/template"
)

type javaClientGenerator struct {
	registry rdl.TypeRegistry
	schema   *rdl.Schema
	name     string
	writer   *bufio.Writer
	err      error
	banner   string
	ns       string
	base     string
}

// GenerateJavaClient generates the client code to talk to the server
func GenerateJavaClient(banner string, schema *rdl.Schema, outdir string, ns string, base string, options []string) error {
	reg := rdl.NewTypeRegistry(schema)
	packageDir, err := javaGenerationDir(outdir, schema, ns)
	if err != nil {
		return err
	}

	cName := javaGenerationStringOptionSet(options, "clientclass")
	if cName == "" {
		cName = capitalize(string(schema.Name))
	}

	out, file, _, err := outputWriter(packageDir, cName, "Client.java")
	if err != nil {
		return err
	}
	gen := &javaClientGenerator{reg, schema, cName, out, nil, banner, ns, base}
	gen.processTemplate(javaClientTemplate)
	out.Flush()
	file.Close()
	if gen.err != nil {
		return gen.err
	}

	//ResourceException - the throawable wrapper for alternate return types
	out, file, _, err = outputWriter(packageDir, "ResourceException", ".java")
	if err != nil {
		return err
	}
	err = javaGenerateResourceException(schema, out, ns)
	out.Flush()
	file.Close()
	if err != nil {
		return err
	}

	//ResourceError - the default data object for an error
	out, file, _, err = outputWriter(packageDir, "ResourceError", ".java")
	if err != nil {
		return err
	}
	err = javaGenerateResourceError(schema, out, ns)
	out.Flush()
	file.Close()
	return err
}

func (gen *javaClientGenerator) processTemplate(templateSource string) error {
	commentFun := func(s string) string {
		return formatComment(s, 0, 80)
	}
	funcMap := template.FuncMap{
		"header":     func() string { return javaGenerationHeader(gen.banner) },
		"package":    func() string { return javaGenerationPackage(gen.schema, gen.ns) },
		"comment":    commentFun,
		"methodSig":  func(r *rdl.Resource) string { return gen.clientMethodSignature(r) },
		"methodBody": func(r *rdl.Resource) string { return gen.clientMethodBody(r) },
		"name":       func() string { return gen.name },
		"cName":      func() string { return capitalize(gen.name) },
		"lName":      func() string { return uncapitalize(gen.name) },
	}
	t := template.Must(template.New(gen.name).Funcs(funcMap).Parse(templateSource))
	return t.Execute(gen.writer, gen.schema)
}

func (gen *javaClientGenerator) resourcePath(r *rdl.Resource) string {
	path := r.Path
	i := strings.Index(path, "?")
	if i >= 0 {
		path = path[0:i]
	}
	return path
}

const javaClientTemplate = `{{header}}
package {{package}};
import com.yahoo.rdl.*;
import javax.ws.rs.client.*;
import javax.ws.rs.*;
import javax.ws.rs.core.*;
import javax.net.ssl.HostnameVerifier;

public class {{cName}}Client {
    Client client;
    WebTarget base;
    String credsHeader;
    String credsToken;

    public {{cName}}Client(String url) {
        client = ClientBuilder.newClient();
        base = client.target(url);
    }

    public {{cName}}Client(String url, HostnameVerifier hostnameVerifier) {
        client = ClientBuilder.newBuilder()
            .hostnameVerifier(hostnameVerifier)
            .build();
        base = client.target(url);
    }

    public void close() {
        client.close();
    }

    public {{cName}}Client setProperty(String name, Object value) {
        client = client.property(name, value);
        return this;
    }

    public {{cName}}Client addCredentials(String header, String token) {
        credsHeader = header;
        credsToken = token;
        return this;
    }
{{range .Resources}}
    {{methodSig .}} {
        {{methodBody .}}
    }
{{end}}
}
`

func (gen *javaClientGenerator) clientMethodSignature(r *rdl.Resource) string {
	reg := gen.registry
	returnType := javaType(reg, r.Type, false, "", "")
	methName, params := javaMethodName(reg, r)
	sparams := ""
	if len(params) > 0 {
		sparams = strings.Join(params, ", ")
	}
	return "public " + returnType + " " + methName + "(" + sparams + ")"
}

func (gen *javaClientGenerator) clientMethodBody(r *rdl.Resource) string {
	reg := gen.registry
	returnType := javaType(reg, r.Type, false, "", "")
	path := r.Path
	s := "WebTarget target = base.path(\"" + path + "\")"
	entityName := ""
	q := ""
	h := ""
	for _, in := range r.Inputs {
		iname := javaName(in.Name)
		if in.PathParam {
			s += "\n            .resolveTemplate(\"" + iname + "\", " + iname + ")"
		} else if in.QueryParam != "" {
			q += "\n        if (" + iname + " != null) {"
			q += "\n            target = target.queryParam(\"" + in.QueryParam + "\", " + iname + ");"
			q += "\n        }"
		} else if in.Header != "" {
			h += "\n        if (" + iname + " != null) {"
			h += "\n            invocationBuilder = invocationBuilder.header(\"" + in.Header + "\", " + iname + ");"
			h += "\n        }"
		} else { //the entity
			entityName = iname
		}
	}
	s += ";"
	if q != "" {
		s += q
	}
	s += "\n        Invocation.Builder invocationBuilder = target.request(\"application/json\");"
	if h != "" {
		s += h
	}
	s += "\n"
	switch r.Method {
	case "PUT", "POST":
		s += "        Response response = invocationBuilder." + strings.ToLower(r.Method) + "(javax.ws.rs.client.Entity.entity(" + entityName + ", \"application/json\"));\n"
	default:
		s += "        Response response = invocationBuilder." + strings.ToLower(r.Method) + "();\n"
	}
	s += "        int code = response.getStatus();\n"
	s += "        String status = Response.Status.fromStatusCode(code).toString();\n"
	expected := "status.equals(\"" + r.Expected + "\")"
	for _, alt := range r.Alternatives {
		expected += "|| status.equals(\"" + alt + "\")"
	}
	s += "        if (" + expected + ") {\n"
	s += "            return response.readEntity(" + returnType + ".class);\n"
	if r.Exceptions != nil {
		for k, v := range r.Exceptions {
			s += "        } else if (status.equals(\"" + k + "\")) {\n"
			s += "            throw new ResourceException(code, response.readEntity(" + v.Type + ".class));\n"
		}
	}
	s += "        } else {\n"
	s += "            throw new ResourceException(code, response.readEntity(Object.class));\n"
	s += "        }\n"
	return s
}
<|endoftext|>


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

package main

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"os/user"
	"path/filepath"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"
	"github.com/keybase/gomounts"
)

var kbfusePath = fuse.OSXFUSEPaths{
	DevicePrefix: "/dev/kbfuse",
	Load:         "/Library/Filesystems/kbfuse.fs/Contents/Resources/load_kbfuse",
	Mount:        "/Library/Filesystems/kbfuse.fs/Contents/Resources/mount_kbfuse",
	DaemonVar:    "MOUNT_KBFUSE_DAEMON_PATH",
}

const (
	mountpointTimeout = 5 * time.Second
	notRunningName    = "KBFS_NOT_RUNNING"
)

type symlink struct {
	link string
}

func (s symlink) Attr(ctx context.Context, a *fuse.Attr) (err error) {
	a.Mode = os.ModeSymlink | a.Mode | 0555
	a.Valid = 0
	return nil
}

func (s symlink) Readlink(ctx context.Context, req *fuse.ReadlinkRequest) (
	link string, err error) {
	return s.link, nil
}

type cacheEntry struct {
	mountpoint string
	time       time.Time
}

type root struct {
	runmodeStr      string
	runmodeStrFancy string

	lock            sync.RWMutex
	mountpointCache map[uint32]cacheEntry

	getMountsLock sync.Mutex
}

func newRoot() *root {
	runmodeStr := "keybase"
	runmodeStrFancy := "Keybase"
	switch os.Getenv("KEYBASE_RUN_MODE") {
	case "staging":
		runmodeStr = "keybase.staging"
		runmodeStrFancy = "KeybaseStaging"
	case "devel":
		runmodeStr = "keybase.devel"
		runmodeStrFancy = "KeybaseDevel"
	}

	return &root{
		runmodeStr:      runmodeStr,
		runmodeStrFancy: runmodeStrFancy,
		mountpointCache: make(map[uint32]cacheEntry),
	}
}

func (r *root) Root() (fs.Node, error) {
	return r, nil
}

func (r *root) Attr(ctx context.Context, attr *fuse.Attr) error {
	attr.Mode = os.ModeDir | 0555
	return nil
}

func (r *root) getCachedMountpoint(uid uint32) string {
	r.lock.RLock()
	defer r.lock.RUnlock()
	entry, ok := r.mountpointCache[uid]
	if !ok {
		return ""
	}
	now := time.Now()
	if now.Sub(entry.time) > mountpointTimeout {
		// Don't bother deleting the entry, since the caller should
		// just overwrite it.
		return ""
	}
	return entry.mountpoint
}

func (r *root) getMountedVolumes() ([]gomounts.Volume, error) {
	r.getMountsLock.Lock()
	defer r.getMountsLock.Unlock()
	return gomounts.GetMountedVolumes()
}

// mountpointMatchesRunmode returns true if `mp` contains `runmode` at
// the end of a component of the path, or followed by a space.
func mountpointMatchesRunmode(mp, runmode string) bool {
	i := strings.Index(mp, runmode)
	if i < 0 {
		return false
	}
	if len(mp) == i+len(runmode) || mp[i+len(runmode)] == '/' ||
		mp[i+len(runmode)] == ' ' {
		return true
	}
	return false
}

func (r *root) findKBFSMount(ctx context.Context) (
	mountpoint string, err error) {
	// Get the UID, and crash intentionally if it's not set, because
	// that means we're not compiled against the correct version of
	// bazil.org/fuse.
	uid := ctx.Value(fs.CtxHeaderUIDKey).(uint32)
	// Don't let the root see anything here; we don't want a symlink
	// loop back to this mount.
	if uid == 0 {
		return "", fuse.ENOENT
	}

	mountpoint = r.getCachedMountpoint(uid)
	if mountpoint != "" {
		return mountpoint, nil
	}

	defer func() {
		if err != nil {
			return
		}
		// Cache the entry if we didn't hit an error.
		r.lock.Lock()
		defer r.lock.Unlock()
		r.mountpointCache[uid] = cacheEntry{
			mountpoint: mountpoint,
			time:       time.Now(),
		}
	}()

	u, err := user.LookupId(strconv.FormatUint(uint64(uid), 10))
	if err != nil {
		return "", err
	}

	vols, err := r.getMountedVolumes()
	if err != nil {
		return "", err
	}
	fuseType := "fuse"
	if runtime.GOOS == "darwin" {
		fuseType = "kbfuse"
	}
	var fuseMountPoints []string
	for _, v := range vols {
		if v.Type != fuseType {
			continue
		}
		if v.Owner != u.Uid {
			continue
		}
		fuseMountPoints = append(fuseMountPoints, v.Path)
	}

	if len(fuseMountPoints) == 0 {
		return "", fuse.ENOENT
	}

	// Pick the first one alphabetically that has "keybase" in the
	// path.
	sort.Strings(fuseMountPoints)
	for _, mp := range fuseMountPoints {
		// Find mountpoints like "/home/user/.local/share/keybase/fs",
		// or "/Volumes/Keybase (user)", and make sure it doesn't
		// match mounts for another run mode, say
		// "/Volumes/KeybaseStaging (user)".
		if mountpointMatchesRunmode(mp, r.runmodeStr) ||
			mountpointMatchesRunmode(mp, r.runmodeStrFancy) {
			return mp, nil
		}
	}

	// Give up.
	return "", fuse.ENOENT
}

func (r *root) ReadDirAll(ctx context.Context) ([]fuse.Dirent, error) {
	_, err := r.findKBFSMount(ctx)
	if err != nil {
		if err == fuse.ENOENT {
			// Put a symlink in the directory for someone who's not
			// logged in, so that the directory is non-empty and
			// future redirector calls as root won't try to mount over
			// us.
			return []fuse.Dirent{
				{
					Type: fuse.DT_Link,
					Name: notRunningName,
				},
			}, nil
		}
		return []fuse.Dirent{}, err
	}

	// TODO: show the `kbfs.error.txt" and "kbfs.nologin.txt" files if
	// they exist?  As root, it is hard to figure out if they're
	// there, though.
	return []fuse.Dirent{
		{
			Type: fuse.DT_Link,
			Name: "private",
		},
		{
			Type: fuse.DT_Link,
			Name: "public",
		},
		fuse.Dirent{
			Type: fuse.DT_Link,
			Name: "team",
		},
	}, nil
}

func (r *root) Lookup(
	ctx context.Context, req *fuse.LookupRequest, resp *fuse.LookupResponse) (
	n fs.Node, err error) {
	mountpoint, err := r.findKBFSMount(ctx)
	if err != nil {
		if req.Name == notRunningName {
			return symlink{"/dev/null"}, nil
		}
		return nil, err
	}

	resp.EntryValid = 0
	switch req.Name {
	case "private", "public", "team", ".kbfs_error", ".kbfs_metrics",
		".kbfs_profiles", ".kbfs_reset_caches", ".kbfs_status",
		"kbfs.error.txt", "kbfs.nologin.txt", ".kbfs_enable_auto_journals",
		".kbfs_disable_auto_journals", ".kbfs_enable_block_prefetching",
		".kbfs_disable_block_prefetching", ".kbfs_enable_debug_server",
		".kbfs_disable_debug_server":
		return symlink{filepath.Join(mountpoint, req.Name)}, nil
	}
	return nil, fuse.ENOENT
}

func main() {
	if len(os.Args) != 2 {
		fmt.Fprintf(os.Stderr, "Usage: %s <mountpoint>\n", os.Args[0])
		os.Exit(1)
	}

	// Restrict the mountpoint to paths starting with "/keybase".
	// Since this is a suid binary, it is dangerous to allow arbitrary
	// mountpoints.  TODO: Read a redirector mountpoint from a
	// root-owned config file.
	r := newRoot()
	if os.Args[1] != fmt.Sprintf("/%s", r.runmodeStr) {
		fmt.Fprintf(os.Stderr, "ERROR: The redirector may only mount at "+
			"/%s; %s is an invalid mountpoint\n", r.runmodeStr, os.Args[1])
		os.Exit(1)
	}

	currUser, err := user.Current()
	if err != nil {
		panic(err)
	}
	if currUser.Uid != "0" {
		runtime.LockOSThread()
		_, _, errNo := syscall.Syscall(syscall.SYS_SETUID, 0, 0, 0)
		if errNo != 0 {
			fmt.Fprintf(os.Stderr, "Can't setuid: %+v\n", errNo)
			os.Exit(1)
		}
	}

	options := []fuse.MountOption{fuse.AllowOther()}
	options = append(options, fuse.FSName("keybase-redirector"))
	options = append(options, fuse.ReadOnly())
	if runtime.GOOS == "darwin" {
		options = append(options, fuse.OSXFUSELocations(kbfusePath))
		options = append(options, fuse.VolumeName("keybase"))
		options = append(options, fuse.NoBrowse())
		// Without NoLocalCaches(), OSX will cache symlinks for a long time.
		options = append(options, fuse.NoLocalCaches())
	}

	c, err := fuse.Mount(os.Args[1], options...)
	if err != nil {
		fmt.Printf("Mount error, exiting cleanly: %+v\n", err)
		os.Exit(0)
	}

	interruptChan := make(chan os.Signal, 1)
	signal.Notify(interruptChan, os.Interrupt)
	signal.Notify(interruptChan, syscall.SIGTERM)
	go func() {
		_ = <-interruptChan

		// This might be a different system thread than above, so we
		// might need to setuid again.
		runtime.LockOSThread()
		_, _, errNo := syscall.Syscall(syscall.SYS_SETUID, 0, 0, 0)
		if errNo != 0 {
			fmt.Fprintf(os.Stderr, "Can't setuid: %+v\n", errNo)
			os.Exit(1)
		}
		err := fuse.Unmount(os.Args[1])
		if err != nil {
			fmt.Printf("Couldn't unmount cleanly: %+v", err)
		}
	}()

	srv := fs.New(c, &fs.Config{
		WithContext: func(ctx context.Context, _ fuse.Request) context.Context {
			return context.Background()
		},
	})
	srv.Serve(r)
}
<commit_msg>redirector: redirect /keybase/.kbfs_edit_history<commit_after>// Copyright 2018 Keybase Inc. All rights reserved.
// Use of this source code is governed by a BSD
// license that can be found in the LICENSE file.

package main

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"os/user"
	"path/filepath"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"bazil.org/fuse"
	"bazil.org/fuse/fs"
	"github.com/keybase/gomounts"
)

var kbfusePath = fuse.OSXFUSEPaths{
	DevicePrefix: "/dev/kbfuse",
	Load:         "/Library/Filesystems/kbfuse.fs/Contents/Resources/load_kbfuse",
	Mount:        "/Library/Filesystems/kbfuse.fs/Contents/Resources/mount_kbfuse",
	DaemonVar:    "MOUNT_KBFUSE_DAEMON_PATH",
}

const (
	mountpointTimeout = 5 * time.Second
	notRunningName    = "KBFS_NOT_RUNNING"
)

type symlink struct {
	link string
}

func (s symlink) Attr(ctx context.Context, a *fuse.Attr) (err error) {
	a.Mode = os.ModeSymlink | a.Mode | 0555
	a.Valid = 0
	return nil
}

func (s symlink) Readlink(ctx context.Context, req *fuse.ReadlinkRequest) (
	link string, err error) {
	return s.link, nil
}

type cacheEntry struct {
	mountpoint string
	time       time.Time
}

type root struct {
	runmodeStr      string
	runmodeStrFancy string

	lock            sync.RWMutex
	mountpointCache map[uint32]cacheEntry

	getMountsLock sync.Mutex
}

func newRoot() *root {
	runmodeStr := "keybase"
	runmodeStrFancy := "Keybase"
	switch os.Getenv("KEYBASE_RUN_MODE") {
	case "staging":
		runmodeStr = "keybase.staging"
		runmodeStrFancy = "KeybaseStaging"
	case "devel":
		runmodeStr = "keybase.devel"
		runmodeStrFancy = "KeybaseDevel"
	}

	return &root{
		runmodeStr:      runmodeStr,
		runmodeStrFancy: runmodeStrFancy,
		mountpointCache: make(map[uint32]cacheEntry),
	}
}

func (r *root) Root() (fs.Node, error) {
	return r, nil
}

func (r *root) Attr(ctx context.Context, attr *fuse.Attr) error {
	attr.Mode = os.ModeDir | 0555
	return nil
}

func (r *root) getCachedMountpoint(uid uint32) string {
	r.lock.RLock()
	defer r.lock.RUnlock()
	entry, ok := r.mountpointCache[uid]
	if !ok {
		return ""
	}
	now := time.Now()
	if now.Sub(entry.time) > mountpointTimeout {
		// Don't bother deleting the entry, since the caller should
		// just overwrite it.
		return ""
	}
	return entry.mountpoint
}

func (r *root) getMountedVolumes() ([]gomounts.Volume, error) {
	r.getMountsLock.Lock()
	defer r.getMountsLock.Unlock()
	return gomounts.GetMountedVolumes()
}

// mountpointMatchesRunmode returns true if `mp` contains `runmode` at
// the end of a component of the path, or followed by a space.
func mountpointMatchesRunmode(mp, runmode string) bool {
	i := strings.Index(mp, runmode)
	if i < 0 {
		return false
	}
	if len(mp) == i+len(runmode) || mp[i+len(runmode)] == '/' ||
		mp[i+len(runmode)] == ' ' {
		return true
	}
	return false
}

func (r *root) findKBFSMount(ctx context.Context) (
	mountpoint string, err error) {
	// Get the UID, and crash intentionally if it's not set, because
	// that means we're not compiled against the correct version of
	// bazil.org/fuse.
	uid := ctx.Value(fs.CtxHeaderUIDKey).(uint32)
	// Don't let the root see anything here; we don't want a symlink
	// loop back to this mount.
	if uid == 0 {
		return "", fuse.ENOENT
	}

	mountpoint = r.getCachedMountpoint(uid)
	if mountpoint != "" {
		return mountpoint, nil
	}

	defer func() {
		if err != nil {
			return
		}
		// Cache the entry if we didn't hit an error.
		r.lock.Lock()
		defer r.lock.Unlock()
		r.mountpointCache[uid] = cacheEntry{
			mountpoint: mountpoint,
			time:       time.Now(),
		}
	}()

	u, err := user.LookupId(strconv.FormatUint(uint64(uid), 10))
	if err != nil {
		return "", err
	}

	vols, err := r.getMountedVolumes()
	if err != nil {
		return "", err
	}
	fuseType := "fuse"
	if runtime.GOOS == "darwin" {
		fuseType = "kbfuse"
	}
	var fuseMountPoints []string
	for _, v := range vols {
		if v.Type != fuseType {
			continue
		}
		if v.Owner != u.Uid {
			continue
		}
		fuseMountPoints = append(fuseMountPoints, v.Path)
	}

	if len(fuseMountPoints) == 0 {
		return "", fuse.ENOENT
	}

	// Pick the first one alphabetically that has "keybase" in the
	// path.
	sort.Strings(fuseMountPoints)
	for _, mp := range fuseMountPoints {
		// Find mountpoints like "/home/user/.local/share/keybase/fs",
		// or "/Volumes/Keybase (user)", and make sure it doesn't
		// match mounts for another run mode, say
		// "/Volumes/KeybaseStaging (user)".
		if mountpointMatchesRunmode(mp, r.runmodeStr) ||
			mountpointMatchesRunmode(mp, r.runmodeStrFancy) {
			return mp, nil
		}
	}

	// Give up.
	return "", fuse.ENOENT
}

func (r *root) ReadDirAll(ctx context.Context) ([]fuse.Dirent, error) {
	_, err := r.findKBFSMount(ctx)
	if err != nil {
		if err == fuse.ENOENT {
			// Put a symlink in the directory for someone who's not
			// logged in, so that the directory is non-empty and
			// future redirector calls as root won't try to mount over
			// us.
			return []fuse.Dirent{
				{
					Type: fuse.DT_Link,
					Name: notRunningName,
				},
			}, nil
		}
		return []fuse.Dirent{}, err
	}

	// TODO: show the `kbfs.error.txt" and "kbfs.nologin.txt" files if
	// they exist?  As root, it is hard to figure out if they're
	// there, though.
	return []fuse.Dirent{
		{
			Type: fuse.DT_Link,
			Name: "private",
		},
		{
			Type: fuse.DT_Link,
			Name: "public",
		},
		fuse.Dirent{
			Type: fuse.DT_Link,
			Name: "team",
		},
	}, nil
}

func (r *root) Lookup(
	ctx context.Context, req *fuse.LookupRequest, resp *fuse.LookupResponse) (
	n fs.Node, err error) {
	mountpoint, err := r.findKBFSMount(ctx)
	if err != nil {
		if req.Name == notRunningName {
			return symlink{"/dev/null"}, nil
		}
		return nil, err
	}

	resp.EntryValid = 0
	switch req.Name {
	case "private", "public", "team", ".kbfs_error", ".kbfs_metrics",
		".kbfs_profiles", ".kbfs_reset_caches", ".kbfs_status",
		"kbfs.error.txt", "kbfs.nologin.txt", ".kbfs_enable_auto_journals",
		".kbfs_disable_auto_journals", ".kbfs_enable_block_prefetching",
		".kbfs_disable_block_prefetching", ".kbfs_enable_debug_server",
		".kbfs_disable_debug_server", ".kbfs_edit_history":
		return symlink{filepath.Join(mountpoint, req.Name)}, nil
	}
	return nil, fuse.ENOENT
}

func main() {
	if len(os.Args) != 2 {
		fmt.Fprintf(os.Stderr, "Usage: %s <mountpoint>\n", os.Args[0])
		os.Exit(1)
	}

	// Restrict the mountpoint to paths starting with "/keybase".
	// Since this is a suid binary, it is dangerous to allow arbitrary
	// mountpoints.  TODO: Read a redirector mountpoint from a
	// root-owned config file.
	r := newRoot()
	if os.Args[1] != fmt.Sprintf("/%s", r.runmodeStr) {
		fmt.Fprintf(os.Stderr, "ERROR: The redirector may only mount at "+
			"/%s; %s is an invalid mountpoint\n", r.runmodeStr, os.Args[1])
		os.Exit(1)
	}

	currUser, err := user.Current()
	if err != nil {
		panic(err)
	}
	if currUser.Uid != "0" {
		runtime.LockOSThread()
		_, _, errNo := syscall.Syscall(syscall.SYS_SETUID, 0, 0, 0)
		if errNo != 0 {
			fmt.Fprintf(os.Stderr, "Can't setuid: %+v\n", errNo)
			os.Exit(1)
		}
	}

	options := []fuse.MountOption{fuse.AllowOther()}
	options = append(options, fuse.FSName("keybase-redirector"))
	options = append(options, fuse.ReadOnly())
	if runtime.GOOS == "darwin" {
		options = append(options, fuse.OSXFUSELocations(kbfusePath))
		options = append(options, fuse.VolumeName("keybase"))
		options = append(options, fuse.NoBrowse())
		// Without NoLocalCaches(), OSX will cache symlinks for a long time.
		options = append(options, fuse.NoLocalCaches())
	}

	c, err := fuse.Mount(os.Args[1], options...)
	if err != nil {
		fmt.Printf("Mount error, exiting cleanly: %+v\n", err)
		os.Exit(0)
	}

	interruptChan := make(chan os.Signal, 1)
	signal.Notify(interruptChan, os.Interrupt)
	signal.Notify(interruptChan, syscall.SIGTERM)
	go func() {
		_ = <-interruptChan

		// This might be a different system thread than above, so we
		// might need to setuid again.
		runtime.LockOSThread()
		_, _, errNo := syscall.Syscall(syscall.SYS_SETUID, 0, 0, 0)
		if errNo != 0 {
			fmt.Fprintf(os.Stderr, "Can't setuid: %+v\n", errNo)
			os.Exit(1)
		}
		err := fuse.Unmount(os.Args[1])
		if err != nil {
			fmt.Printf("Couldn't unmount cleanly: %+v", err)
		}
	}()

	srv := fs.New(c, &fs.Config{
		WithContext: func(ctx context.Context, _ fuse.Request) context.Context {
			return context.Background()
		},
	})
	srv.Serve(r)
}
<|endoftext|>