github.com/google/go-safeweb@v0.0.0-20231219055052-64d8cfc90fbb/examples/sample-application/storage/fakedb.go (about)

     1  // Copyright 2020 Google LLC
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //	https://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  //go:build go1.16
    16  // +build go1.16
    17  
    18  package storage
    19  
    20  import (
    21  	"crypto/rand"
    22  	"encoding/base64"
    23  	"errors"
    24  	"sync"
    25  
    26  	"golang.org/x/crypto/scrypt"
    27  )
    28  
    29  // Note: a real program would connect to a real DB using
    30  // github.com/google/go-safeweb/safesql. This is just a simple storage
    31  // implementation to demonstrate the web framework.
    32  
    33  // Please ignore the contents of this file.
    34  
    35  type Note struct {
    36  	Title, Text string
    37  }
    38  
    39  type DB struct {
    40  	mu sync.Mutex
    41  	// user -> note title -> notes
    42  	notes map[string]map[string]Note
    43  
    44  	// user -> token
    45  	sessionTokens map[string]string
    46  	// token -> user
    47  	userSessions map[string]string
    48  
    49  	// user -> pw hash
    50  	credentials map[string]string
    51  }
    52  
    53  func NewDB() *DB {
    54  	return &DB{
    55  		notes:         map[string]map[string]Note{},
    56  		sessionTokens: map[string]string{},
    57  		userSessions:  map[string]string{},
    58  		credentials:   map[string]string{},
    59  	}
    60  }
    61  
    62  // Notes
    63  
    64  func (s *DB) AddOrEditNote(user string, n Note) {
    65  	s.mu.Lock()
    66  	defer s.mu.Unlock()
    67  	if s.notes[user] == nil {
    68  		s.notes[user] = map[string]Note{}
    69  	}
    70  	s.notes[user][n.Title] = n
    71  }
    72  
    73  func (s *DB) GetNotes(user string) []Note {
    74  	s.mu.Lock()
    75  	defer s.mu.Unlock()
    76  	var ns []Note
    77  	for _, n := range s.notes[user] {
    78  		ns = append(ns, n)
    79  	}
    80  	return ns
    81  }
    82  
    83  // Sessions
    84  
    85  func (s *DB) GetUser(token string) (user string, valid bool) {
    86  	s.mu.Lock()
    87  	defer s.mu.Unlock()
    88  	user, valid = s.sessionTokens[token]
    89  	return user, valid
    90  }
    91  
    92  func (s *DB) GetToken(user string) (token string) {
    93  	s.mu.Lock()
    94  	defer s.mu.Unlock()
    95  	token, has := s.userSessions[user]
    96  	if has {
    97  		return token
    98  	}
    99  	token = genToken()
   100  	s.userSessions[user] = token
   101  	s.sessionTokens[token] = user
   102  	return token
   103  }
   104  
   105  func (s *DB) DelSession(user string) {
   106  	s.mu.Lock()
   107  	defer s.mu.Unlock()
   108  	token, has := s.userSessions[user]
   109  	if !has {
   110  		return
   111  	}
   112  	delete(s.userSessions, user)
   113  	delete(s.sessionTokens, token)
   114  }
   115  
   116  func genToken() string {
   117  	b := make([]byte, 20)
   118  	rand.Read(b)
   119  	tok := base64.RawStdEncoding.EncodeToString(b)
   120  	return tok
   121  }
   122  
   123  // Credentials
   124  
   125  // HasUser checks if the user exists.
   126  func (s *DB) HasUser(name string) bool {
   127  	s.mu.Lock()
   128  	defer s.mu.Unlock()
   129  	_, has := s.credentials[name]
   130  	return has
   131  }
   132  
   133  // AddUser adds a user to the storage if it is not already there.
   134  func (s *DB) AddOrAuthUser(name, password string) error {
   135  	s.mu.Lock()
   136  	defer s.mu.Unlock()
   137  	if password == "" {
   138  		return errors.New("password cannot be empty")
   139  	}
   140  	if storedHash, has := s.credentials[name]; has {
   141  		if storedHash != hash(password) {
   142  			return errors.New("wrong password")
   143  		}
   144  		return nil
   145  	}
   146  	s.credentials[name] = hash(password)
   147  	return nil
   148  }
   149  
   150  func hash(pw string) string {
   151  	salt := []byte("please use a proper salt in production")
   152  	hash, err := scrypt.Key([]byte(pw), salt, 32768, 8, 1, 32)
   153  	if err != nil {
   154  		panic("this should not happen")
   155  	}
   156  	return string(hash)
   157  }