package allow import ( "bytes" "encoding/json" "errors" "fmt" "io/ioutil" "os" "path" "sort" "strings" ) const ( AL_QUERY int = iota + 1 AL_VARS ) type Item struct { Name string key string URI string Query string Vars json.RawMessage } type List struct { filepath string saveChan chan Item } type Config struct { CreateIfNotExists bool Persist bool } func New(cpath string, conf Config) (*List, error) { al := List{} if len(cpath) != 0 { fp := path.Join(cpath, "allow.list") if _, err := os.Stat(fp); err == nil { al.filepath = fp } else if !os.IsNotExist(err) { return nil, err } } if len(al.filepath) == 0 { fp := "./allow.list" if _, err := os.Stat(fp); err == nil { al.filepath = fp } else if !os.IsNotExist(err) { return nil, err } } if len(al.filepath) == 0 { fp := "./config/allow.list" if _, err := os.Stat(fp); err == nil { al.filepath = fp } else if !os.IsNotExist(err) { return nil, err } } if len(al.filepath) == 0 { if !conf.CreateIfNotExists { return nil, errors.New("allow.list not found") } if len(cpath) == 0 { al.filepath = "./config/allow.list" } else { al.filepath = path.Join(cpath, "allow.list") } } var err error if conf.Persist { al.saveChan = make(chan Item) go func() { for v := range al.saveChan { if err = al.save(v); err != nil { break } } }() } if err != nil { return nil, err } return &al, nil } func (al *List) IsPersist() bool { return al.saveChan != nil } func (al *List) Add(vars []byte, query, uri string) error { if al.saveChan == nil { return errors.New("allow.list is read-only") } if len(query) == 0 { return errors.New("empty query") } var q string for i := 0; i < len(query); i++ { c := query[i] if c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' { q = query break } else if c == '{' { q = "query " + query break } } al.saveChan <- Item{ URI: uri, Query: q, Vars: vars, } return nil } func (al *List) Load() ([]Item, error) { var list []Item b, err := ioutil.ReadFile(al.filepath) if err != nil { return list, err } if len(b) == 0 { return list, nil } var uri string var varBytes []byte itemMap := make(map[string]struct{}) s, e, c := 0, 0, 0 ty := 0 for { fq := false if c == 0 && b[e] == '#' { s = e for e < len(b) && b[e] != '\n' { e++ } if (e - s) > 2 { uri = strings.TrimSpace(string(b[(s + 1):e])) } } if e >= len(b) { break } if matchPrefix(b, e, "query") || matchPrefix(b, e, "mutation") { if c == 0 { s = e } ty = AL_QUERY } else if matchPrefix(b, e, "variables") { if c == 0 { s = e + len("variables") + 1 } ty = AL_VARS } else if b[e] == '{' { c++ } else if b[e] == '}' { c-- if c == 0 { if ty == AL_QUERY { fq = true } else if ty == AL_VARS { varBytes = b[s:(e + 1)] } ty = 0 } } if fq { query := string(b[s:(e + 1)]) name := QueryName(query) key := strings.ToLower(name) if _, ok := itemMap[key]; !ok { v := Item{ Name: name, key: key, URI: uri, Query: query, Vars: varBytes, } list = append(list, v) } varBytes = nil } e++ if e >= len(b) { break } } return list, nil } func (al *List) save(item Item) error { item.Name = QueryName(item.Query) item.key = strings.ToLower(item.Name) if len(item.Name) == 0 { return nil } list, err := al.Load() if err != nil { return err } index := -1 for i, v := range list { if strings.EqualFold(v.Name, item.Name) { index = i break } } if index != -1 { list[index] = item } else { list = append(list, item) } f, err := os.Create(al.filepath) if err != nil { return err } defer f.Close() sort.Slice(list, func(i, j int) bool { return strings.Compare(list[i].key, list[j].key) == -1 }) for _, v := range list { _, err := f.WriteString(fmt.Sprintf("# %s\n\n", v.URI)) if err != nil { return err } if len(v.Vars) != 0 && !bytes.Equal(v.Vars, []byte("{}")) { vj, err := json.MarshalIndent(v.Vars, "", " ") if err != nil { return fmt.Errorf("failed to marshal vars: %v", err) } _, err = f.WriteString(fmt.Sprintf("variables %s\n\n", vj)) if err != nil { return err } } if v.Query[0] == '{' { _, err = f.WriteString(fmt.Sprintf("query %s\n\n", v.Query)) } else { _, err = f.WriteString(fmt.Sprintf("%s\n\n", v.Query)) } if err != nil { return err } } return nil } func matchPrefix(b []byte, i int, s string) bool { if (len(b) - i) < len(s) { return false } for n := 0; n < len(s); n++ { if b[(i+n)] != s[n] { return false } } return true } func QueryName(b string) string { state, s := 0, 0 for i := 0; i < len(b); i++ { switch { case state == 2 && b[i] == '{': return b[s:i] case state == 2 && b[i] == ' ': return b[s:i] case state == 1 && b[i] == '{': return "" case state == 1 && b[i] != ' ': s = i state = 2 case state == 1 && b[i] == ' ': continue case i != 0 && b[i] == ' ' && (b[i-1] == 'n' || b[i-1] == 'y'): state = 1 } } return "" }