images.go 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032
  1. package server
  2. import (
  3. "bytes"
  4. "crypto/sha256"
  5. "encoding/json"
  6. "errors"
  7. "fmt"
  8. "io"
  9. "io/ioutil"
  10. "log"
  11. "net/http"
  12. "os"
  13. "path"
  14. "path/filepath"
  15. "reflect"
  16. "strconv"
  17. "strings"
  18. "text/template"
  19. "github.com/jmorganca/ollama/api"
  20. "github.com/jmorganca/ollama/parser"
  21. )
  22. type RegistryOptions struct {
  23. Insecure bool
  24. Username string
  25. Password string
  26. }
  27. type Model struct {
  28. Name string `json:"name"`
  29. ModelPath string
  30. Template string
  31. System string
  32. Options api.Options
  33. }
  34. func (m *Model) Prompt(request api.GenerateRequest) (string, error) {
  35. tmpl, err := template.New("").Parse(m.Template)
  36. if err != nil {
  37. return "", err
  38. }
  39. var vars struct {
  40. First bool
  41. System string
  42. Prompt string
  43. // deprecated: versions <= 0.0.7 used this to omit the system prompt
  44. Context []int
  45. }
  46. vars.First = len(vars.Context) == 0
  47. vars.System = m.System
  48. vars.Prompt = request.Prompt
  49. vars.Context = request.Context
  50. var sb strings.Builder
  51. if err := tmpl.Execute(&sb, vars); err != nil {
  52. return "", err
  53. }
  54. return sb.String(), nil
  55. }
  56. type ManifestV2 struct {
  57. SchemaVersion int `json:"schemaVersion"`
  58. MediaType string `json:"mediaType"`
  59. Config Layer `json:"config"`
  60. Layers []*Layer `json:"layers"`
  61. }
  62. type Layer struct {
  63. MediaType string `json:"mediaType"`
  64. Digest string `json:"digest"`
  65. Size int `json:"size"`
  66. }
  67. type LayerReader struct {
  68. Layer
  69. io.Reader
  70. }
  71. type ConfigV2 struct {
  72. Architecture string `json:"architecture"`
  73. OS string `json:"os"`
  74. RootFS RootFS `json:"rootfs"`
  75. }
  76. type RootFS struct {
  77. Type string `json:"type"`
  78. DiffIDs []string `json:"diff_ids"`
  79. }
  80. func (m *ManifestV2) GetTotalSize() int {
  81. var total int
  82. for _, layer := range m.Layers {
  83. total += layer.Size
  84. }
  85. total += m.Config.Size
  86. return total
  87. }
  88. func GetManifest(mp ModelPath) (*ManifestV2, error) {
  89. fp, err := mp.GetManifestPath(false)
  90. if err != nil {
  91. return nil, err
  92. }
  93. if _, err = os.Stat(fp); err != nil && !errors.Is(err, os.ErrNotExist) {
  94. return nil, fmt.Errorf("couldn't find model '%s'", mp.GetShortTagname())
  95. }
  96. var manifest *ManifestV2
  97. bts, err := os.ReadFile(fp)
  98. if err != nil {
  99. return nil, fmt.Errorf("couldn't open file '%s'", fp)
  100. }
  101. if err := json.Unmarshal(bts, &manifest); err != nil {
  102. return nil, err
  103. }
  104. return manifest, nil
  105. }
  106. func GetModel(name string) (*Model, error) {
  107. mp := ParseModelPath(name)
  108. manifest, err := GetManifest(mp)
  109. if err != nil {
  110. return nil, err
  111. }
  112. model := &Model{
  113. Name: mp.GetFullTagname(),
  114. }
  115. for _, layer := range manifest.Layers {
  116. filename, err := GetBlobsPath(layer.Digest)
  117. if err != nil {
  118. return nil, err
  119. }
  120. switch layer.MediaType {
  121. case "application/vnd.ollama.image.model":
  122. model.ModelPath = filename
  123. case "application/vnd.ollama.image.template":
  124. bts, err := os.ReadFile(filename)
  125. if err != nil {
  126. return nil, err
  127. }
  128. model.Template = string(bts)
  129. case "application/vnd.ollama.image.system":
  130. bts, err := os.ReadFile(filename)
  131. if err != nil {
  132. return nil, err
  133. }
  134. model.System = string(bts)
  135. case "application/vnd.ollama.image.prompt":
  136. bts, err := os.ReadFile(filename)
  137. if err != nil {
  138. return nil, err
  139. }
  140. model.Template = string(bts)
  141. case "application/vnd.ollama.image.params":
  142. params, err := os.Open(filename)
  143. if err != nil {
  144. return nil, err
  145. }
  146. defer params.Close()
  147. var opts api.Options
  148. if err = json.NewDecoder(params).Decode(&opts); err != nil {
  149. return nil, err
  150. }
  151. model.Options = opts
  152. }
  153. }
  154. return model, nil
  155. }
  156. func CreateModel(name string, path string, fn func(status string)) error {
  157. mf, err := os.Open(path)
  158. if err != nil {
  159. fn(fmt.Sprintf("couldn't open modelfile '%s'", path))
  160. return fmt.Errorf("failed to open file: %w", err)
  161. }
  162. defer mf.Close()
  163. fn("parsing modelfile")
  164. commands, err := parser.Parse(mf)
  165. if err != nil {
  166. return err
  167. }
  168. var layers []*LayerReader
  169. params := make(map[string]string)
  170. for _, c := range commands {
  171. log.Printf("[%s] - %s\n", c.Name, c.Args)
  172. switch c.Name {
  173. case "model":
  174. fn("looking for model")
  175. mf, err := GetManifest(ParseModelPath(c.Args))
  176. if err != nil {
  177. fp := c.Args
  178. // If filePath starts with ~/, replace it with the user's home directory.
  179. if strings.HasPrefix(fp, "~/") {
  180. parts := strings.Split(fp, "/")
  181. home, err := os.UserHomeDir()
  182. if err != nil {
  183. return fmt.Errorf("failed to open file: %v", err)
  184. }
  185. fp = filepath.Join(home, filepath.Join(parts[1:]...))
  186. }
  187. // If filePath is not an absolute path, make it relative to the modelfile path
  188. if !filepath.IsAbs(fp) {
  189. fp = filepath.Join(filepath.Dir(path), fp)
  190. }
  191. fn("creating model layer")
  192. file, err := os.Open(fp)
  193. if err != nil {
  194. return fmt.Errorf("failed to open file: %v", err)
  195. }
  196. defer file.Close()
  197. l, err := CreateLayer(file)
  198. if err != nil {
  199. return fmt.Errorf("failed to create layer: %v", err)
  200. }
  201. l.MediaType = "application/vnd.ollama.image.model"
  202. layers = append(layers, l)
  203. } else {
  204. log.Printf("manifest = %#v", mf)
  205. for _, l := range mf.Layers {
  206. newLayer, err := GetLayerWithBufferFromLayer(l)
  207. if err != nil {
  208. return err
  209. }
  210. layers = append(layers, newLayer)
  211. }
  212. }
  213. case "license", "template", "system", "prompt":
  214. fn(fmt.Sprintf("creating %s layer", c.Name))
  215. // remove the prompt layer if one exists
  216. mediaType := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)
  217. layers = removeLayerFromLayers(layers, mediaType)
  218. layer, err := CreateLayer(strings.NewReader(c.Args))
  219. if err != nil {
  220. return err
  221. }
  222. layer.MediaType = mediaType
  223. layers = append(layers, layer)
  224. default:
  225. params[c.Name] = c.Args
  226. }
  227. }
  228. // Create a single layer for the parameters
  229. if len(params) > 0 {
  230. fn("creating parameter layer")
  231. layers = removeLayerFromLayers(layers, "application/vnd.ollama.image.params")
  232. paramData, err := paramsToReader(params)
  233. if err != nil {
  234. return fmt.Errorf("couldn't create params json: %v", err)
  235. }
  236. l, err := CreateLayer(paramData)
  237. if err != nil {
  238. return fmt.Errorf("failed to create layer: %v", err)
  239. }
  240. l.MediaType = "application/vnd.ollama.image.params"
  241. layers = append(layers, l)
  242. }
  243. digests, err := getLayerDigests(layers)
  244. if err != nil {
  245. return err
  246. }
  247. var manifestLayers []*Layer
  248. for _, l := range layers {
  249. manifestLayers = append(manifestLayers, &l.Layer)
  250. }
  251. // Create a layer for the config object
  252. fn("creating config layer")
  253. cfg, err := createConfigLayer(digests)
  254. if err != nil {
  255. return err
  256. }
  257. layers = append(layers, cfg)
  258. err = SaveLayers(layers, fn, false)
  259. if err != nil {
  260. return err
  261. }
  262. // Create the manifest
  263. fn("writing manifest")
  264. err = CreateManifest(name, cfg, manifestLayers)
  265. if err != nil {
  266. return err
  267. }
  268. fn("success")
  269. return nil
  270. }
  271. func removeLayerFromLayers(layers []*LayerReader, mediaType string) []*LayerReader {
  272. j := 0
  273. for _, l := range layers {
  274. if l.MediaType != mediaType {
  275. layers[j] = l
  276. j++
  277. }
  278. }
  279. return layers[:j]
  280. }
  281. func SaveLayers(layers []*LayerReader, fn func(status string), force bool) error {
  282. // Write each of the layers to disk
  283. for _, layer := range layers {
  284. fp, err := GetBlobsPath(layer.Digest)
  285. if err != nil {
  286. return err
  287. }
  288. _, err = os.Stat(fp)
  289. if os.IsNotExist(err) || force {
  290. fn(fmt.Sprintf("writing layer %s", layer.Digest))
  291. out, err := os.Create(fp)
  292. if err != nil {
  293. log.Printf("couldn't create %s", fp)
  294. return err
  295. }
  296. defer out.Close()
  297. if _, err = io.Copy(out, layer.Reader); err != nil {
  298. return err
  299. }
  300. } else {
  301. fn(fmt.Sprintf("using already created layer %s", layer.Digest))
  302. }
  303. }
  304. return nil
  305. }
  306. func CreateManifest(name string, cfg *LayerReader, layers []*Layer) error {
  307. mp := ParseModelPath(name)
  308. manifest := ManifestV2{
  309. SchemaVersion: 2,
  310. MediaType: "application/vnd.docker.distribution.manifest.v2+json",
  311. Config: Layer{
  312. MediaType: cfg.MediaType,
  313. Size: cfg.Size,
  314. Digest: cfg.Digest,
  315. },
  316. Layers: layers,
  317. }
  318. manifestJSON, err := json.Marshal(manifest)
  319. if err != nil {
  320. return err
  321. }
  322. fp, err := mp.GetManifestPath(true)
  323. if err != nil {
  324. return err
  325. }
  326. return os.WriteFile(fp, manifestJSON, 0o644)
  327. }
  328. func GetLayerWithBufferFromLayer(layer *Layer) (*LayerReader, error) {
  329. fp, err := GetBlobsPath(layer.Digest)
  330. if err != nil {
  331. return nil, err
  332. }
  333. file, err := os.Open(fp)
  334. if err != nil {
  335. return nil, fmt.Errorf("could not open blob: %w", err)
  336. }
  337. defer file.Close()
  338. newLayer, err := CreateLayer(file)
  339. if err != nil {
  340. return nil, err
  341. }
  342. newLayer.MediaType = layer.MediaType
  343. return newLayer, nil
  344. }
  345. func paramsToReader(params map[string]string) (io.ReadSeeker, error) {
  346. opts := api.DefaultOptions()
  347. typeOpts := reflect.TypeOf(opts)
  348. // build map of json struct tags
  349. jsonOpts := make(map[string]reflect.StructField)
  350. for _, field := range reflect.VisibleFields(typeOpts) {
  351. jsonTag := strings.Split(field.Tag.Get("json"), ",")[0]
  352. if jsonTag != "" {
  353. jsonOpts[jsonTag] = field
  354. }
  355. }
  356. valueOpts := reflect.ValueOf(&opts).Elem()
  357. // iterate params and set values based on json struct tags
  358. for key, val := range params {
  359. if opt, ok := jsonOpts[key]; ok {
  360. field := valueOpts.FieldByName(opt.Name)
  361. if field.IsValid() && field.CanSet() {
  362. switch field.Kind() {
  363. case reflect.Float32:
  364. floatVal, err := strconv.ParseFloat(val, 32)
  365. if err != nil {
  366. return nil, fmt.Errorf("invalid float value %s", val)
  367. }
  368. field.SetFloat(floatVal)
  369. case reflect.Int:
  370. intVal, err := strconv.ParseInt(val, 10, 0)
  371. if err != nil {
  372. return nil, fmt.Errorf("invalid int value %s", val)
  373. }
  374. field.SetInt(intVal)
  375. case reflect.Bool:
  376. boolVal, err := strconv.ParseBool(val)
  377. if err != nil {
  378. return nil, fmt.Errorf("invalid bool value %s", val)
  379. }
  380. field.SetBool(boolVal)
  381. case reflect.String:
  382. field.SetString(val)
  383. default:
  384. return nil, fmt.Errorf("unknown type %s for %s", field.Kind(), key)
  385. }
  386. }
  387. }
  388. }
  389. bts, err := json.Marshal(opts)
  390. if err != nil {
  391. return nil, err
  392. }
  393. return bytes.NewReader(bts), nil
  394. }
  395. func getLayerDigests(layers []*LayerReader) ([]string, error) {
  396. var digests []string
  397. for _, l := range layers {
  398. if l.Digest == "" {
  399. return nil, fmt.Errorf("layer is missing a digest")
  400. }
  401. digests = append(digests, l.Digest)
  402. }
  403. return digests, nil
  404. }
  405. // CreateLayer creates a Layer object from a given file
  406. func CreateLayer(f io.ReadSeeker) (*LayerReader, error) {
  407. digest, size := GetSHA256Digest(f)
  408. f.Seek(0, 0)
  409. layer := &LayerReader{
  410. Layer: Layer{
  411. MediaType: "application/vnd.docker.image.rootfs.diff.tar",
  412. Digest: digest,
  413. Size: size,
  414. },
  415. Reader: f,
  416. }
  417. return layer, nil
  418. }
  419. func DeleteModel(name string, fn func(api.ProgressResponse)) error {
  420. mp := ParseModelPath(name)
  421. manifest, err := GetManifest(mp)
  422. if err != nil {
  423. fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
  424. return err
  425. }
  426. deleteMap := make(map[string]bool)
  427. for _, layer := range manifest.Layers {
  428. deleteMap[layer.Digest] = true
  429. }
  430. deleteMap[manifest.Config.Digest] = true
  431. fp, err := GetManifestPath()
  432. if err != nil {
  433. fn(api.ProgressResponse{Status: "problem getting manifest path"})
  434. return err
  435. }
  436. err = filepath.Walk(fp, func(path string, info os.FileInfo, err error) error {
  437. if err != nil {
  438. fn(api.ProgressResponse{Status: "problem walking manifest dir"})
  439. return err
  440. }
  441. if !info.IsDir() {
  442. path := path[len(fp)+1:]
  443. slashIndex := strings.LastIndex(path, "/")
  444. if slashIndex == -1 {
  445. return nil
  446. }
  447. tag := path[:slashIndex] + ":" + path[slashIndex+1:]
  448. fmp := ParseModelPath(tag)
  449. // skip the manifest we're trying to delete
  450. if mp.GetFullTagname() == fmp.GetFullTagname() {
  451. return nil
  452. }
  453. // save (i.e. delete from the deleteMap) any files used in other manifests
  454. manifest, err := GetManifest(fmp)
  455. if err != nil {
  456. log.Printf("skipping file: %s", fp)
  457. return nil
  458. }
  459. for _, layer := range manifest.Layers {
  460. delete(deleteMap, layer.Digest)
  461. }
  462. delete(deleteMap, manifest.Config.Digest)
  463. }
  464. return nil
  465. })
  466. // only delete the files which are still in the deleteMap
  467. for k, v := range deleteMap {
  468. if v {
  469. fp, err := GetBlobsPath(k)
  470. if err != nil {
  471. log.Printf("couldn't get file path for '%s': %v", k, err)
  472. continue
  473. }
  474. if err := os.Remove(fp); err != nil {
  475. log.Printf("couldn't remove file '%s': %v", fp, err)
  476. continue
  477. }
  478. }
  479. }
  480. fp, err = mp.GetManifestPath(false)
  481. if err != nil {
  482. return err
  483. }
  484. err = os.Remove(fp)
  485. if err != nil {
  486. log.Printf("couldn't remove manifest file '%s': %v", fp, err)
  487. return err
  488. }
  489. fn(api.ProgressResponse{Status: fmt.Sprintf("deleted '%s'", name)})
  490. return nil
  491. }
  492. func PushModel(name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
  493. mp := ParseModelPath(name)
  494. fn(api.ProgressResponse{Status: "retrieving manifest"})
  495. manifest, err := GetManifest(mp)
  496. if err != nil {
  497. fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
  498. return err
  499. }
  500. var layers []*Layer
  501. var total int
  502. var completed int
  503. for _, layer := range manifest.Layers {
  504. layers = append(layers, layer)
  505. total += layer.Size
  506. }
  507. layers = append(layers, &manifest.Config)
  508. total += manifest.Config.Size
  509. for _, layer := range layers {
  510. exists, err := checkBlobExistence(mp, layer.Digest, regOpts)
  511. if err != nil {
  512. return err
  513. }
  514. if exists {
  515. completed += layer.Size
  516. fn(api.ProgressResponse{
  517. Status: "using existing layer",
  518. Digest: layer.Digest,
  519. Total: total,
  520. Completed: completed,
  521. })
  522. continue
  523. }
  524. fn(api.ProgressResponse{
  525. Status: "starting upload",
  526. Digest: layer.Digest,
  527. Total: total,
  528. Completed: completed,
  529. })
  530. location, err := startUpload(mp, regOpts)
  531. if err != nil {
  532. log.Printf("couldn't start upload: %v", err)
  533. return err
  534. }
  535. err = uploadBlob(location, layer, regOpts)
  536. if err != nil {
  537. log.Printf("error uploading blob: %v", err)
  538. return err
  539. }
  540. completed += layer.Size
  541. fn(api.ProgressResponse{
  542. Status: "upload complete",
  543. Digest: layer.Digest,
  544. Total: total,
  545. Completed: completed,
  546. })
  547. }
  548. fn(api.ProgressResponse{
  549. Status: "pushing manifest",
  550. Total: total,
  551. Completed: completed,
  552. })
  553. url := fmt.Sprintf("%s/v2/%s/manifests/%s", mp.Registry, mp.GetNamespaceRepository(), mp.Tag)
  554. headers := map[string]string{
  555. "Content-Type": "application/vnd.docker.distribution.manifest.v2+json",
  556. }
  557. manifestJSON, err := json.Marshal(manifest)
  558. if err != nil {
  559. return err
  560. }
  561. resp, err := makeRequest("PUT", url, headers, bytes.NewReader(manifestJSON), regOpts)
  562. if err != nil {
  563. return err
  564. }
  565. defer resp.Body.Close()
  566. // Check for success: For a successful upload, the Docker registry will respond with a 201 Created
  567. if resp.StatusCode != http.StatusCreated {
  568. body, _ := io.ReadAll(resp.Body)
  569. return fmt.Errorf("registry responded with code %d: %v", resp.StatusCode, string(body))
  570. }
  571. fn(api.ProgressResponse{
  572. Status: "success",
  573. Total: total,
  574. Completed: completed,
  575. })
  576. return nil
  577. }
  578. func PullModel(name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
  579. mp := ParseModelPath(name)
  580. fn(api.ProgressResponse{Status: "pulling manifest"})
  581. manifest, err := pullModelManifest(mp, regOpts)
  582. if err != nil {
  583. return fmt.Errorf("pull model manifest: %q", err)
  584. }
  585. var layers []*Layer
  586. layers = append(layers, manifest.Layers...)
  587. layers = append(layers, &manifest.Config)
  588. for _, layer := range layers {
  589. if err := downloadBlob(mp, layer.Digest, regOpts, fn); err != nil {
  590. return err
  591. }
  592. }
  593. fn(api.ProgressResponse{Status: "verifying sha256 digest"})
  594. for _, layer := range layers {
  595. if err := verifyBlob(layer.Digest); err != nil {
  596. return err
  597. }
  598. }
  599. fn(api.ProgressResponse{Status: "writing manifest"})
  600. manifestJSON, err := json.Marshal(manifest)
  601. if err != nil {
  602. return err
  603. }
  604. fp, err := mp.GetManifestPath(true)
  605. if err != nil {
  606. return err
  607. }
  608. err = os.WriteFile(fp, manifestJSON, 0o644)
  609. if err != nil {
  610. log.Printf("couldn't write to %s", fp)
  611. return err
  612. }
  613. fn(api.ProgressResponse{Status: "success"})
  614. return nil
  615. }
  616. func pullModelManifest(mp ModelPath, regOpts *RegistryOptions) (*ManifestV2, error) {
  617. url := fmt.Sprintf("%s/v2/%s/manifests/%s", mp.Registry, mp.GetNamespaceRepository(), mp.Tag)
  618. headers := map[string]string{
  619. "Accept": "application/vnd.docker.distribution.manifest.v2+json",
  620. }
  621. resp, err := makeRequest("GET", url, headers, nil, regOpts)
  622. if err != nil {
  623. log.Printf("couldn't get manifest: %v", err)
  624. return nil, err
  625. }
  626. defer resp.Body.Close()
  627. // Check for success: For a successful upload, the Docker registry will respond with a 201 Created
  628. if resp.StatusCode != http.StatusOK {
  629. body, _ := io.ReadAll(resp.Body)
  630. return nil, fmt.Errorf("registry responded with code %d: %s", resp.StatusCode, body)
  631. }
  632. var m *ManifestV2
  633. if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
  634. return nil, err
  635. }
  636. return m, err
  637. }
  638. func createConfigLayer(layers []string) (*LayerReader, error) {
  639. // TODO change architecture and OS
  640. config := ConfigV2{
  641. Architecture: "arm64",
  642. OS: "linux",
  643. RootFS: RootFS{
  644. Type: "layers",
  645. DiffIDs: layers,
  646. },
  647. }
  648. configJSON, err := json.Marshal(config)
  649. if err != nil {
  650. return nil, err
  651. }
  652. digest, size := GetSHA256Digest(bytes.NewBuffer(configJSON))
  653. layer := &LayerReader{
  654. Layer: Layer{
  655. MediaType: "application/vnd.docker.container.image.v1+json",
  656. Digest: digest,
  657. Size: size,
  658. },
  659. Reader: bytes.NewBuffer(configJSON),
  660. }
  661. return layer, nil
  662. }
  663. // GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer
  664. func GetSHA256Digest(r io.Reader) (string, int) {
  665. h := sha256.New()
  666. n, err := io.Copy(h, r)
  667. if err != nil {
  668. log.Fatal(err)
  669. }
  670. return fmt.Sprintf("sha256:%x", h.Sum(nil)), int(n)
  671. }
  672. func startUpload(mp ModelPath, regOpts *RegistryOptions) (string, error) {
  673. url := fmt.Sprintf("%s/v2/%s/blobs/uploads/", mp.Registry, mp.GetNamespaceRepository())
  674. resp, err := makeRequest("POST", url, nil, nil, regOpts)
  675. if err != nil {
  676. log.Printf("couldn't start upload: %v", err)
  677. return "", err
  678. }
  679. defer resp.Body.Close()
  680. // Check for success
  681. if resp.StatusCode != http.StatusAccepted {
  682. body, _ := io.ReadAll(resp.Body)
  683. return "", fmt.Errorf("registry responded with code %d: %s", resp.StatusCode, body)
  684. }
  685. // Extract UUID location from header
  686. location := resp.Header.Get("Location")
  687. if location == "" {
  688. return "", fmt.Errorf("location header is missing in response")
  689. }
  690. return location, nil
  691. }
  692. // Function to check if a blob already exists in the Docker registry
  693. func checkBlobExistence(mp ModelPath, digest string, regOpts *RegistryOptions) (bool, error) {
  694. url := fmt.Sprintf("%s/v2/%s/blobs/%s", mp.Registry, mp.GetNamespaceRepository(), digest)
  695. resp, err := makeRequest("HEAD", url, nil, nil, regOpts)
  696. if err != nil {
  697. log.Printf("couldn't check for blob: %v", err)
  698. return false, err
  699. }
  700. defer resp.Body.Close()
  701. // Check for success: If the blob exists, the Docker registry will respond with a 200 OK
  702. return resp.StatusCode == http.StatusOK, nil
  703. }
  704. func uploadBlob(location string, layer *Layer, regOpts *RegistryOptions) error {
  705. // Create URL
  706. url := fmt.Sprintf("%s&digest=%s", location, layer.Digest)
  707. headers := make(map[string]string)
  708. headers["Content-Length"] = fmt.Sprintf("%d", layer.Size)
  709. headers["Content-Type"] = "application/octet-stream"
  710. // TODO change from monolithic uploads to chunked uploads
  711. // TODO allow resumability
  712. // TODO allow canceling uploads via DELETE
  713. // TODO allow cross repo blob mount
  714. fp, err := GetBlobsPath(layer.Digest)
  715. if err != nil {
  716. return err
  717. }
  718. f, err := os.Open(fp)
  719. if err != nil {
  720. return err
  721. }
  722. resp, err := makeRequest("PUT", url, headers, f, regOpts)
  723. if err != nil {
  724. log.Printf("couldn't upload blob: %v", err)
  725. return err
  726. }
  727. defer resp.Body.Close()
  728. // Check for success: For a successful upload, the Docker registry will respond with a 201 Created
  729. if resp.StatusCode != http.StatusCreated {
  730. body, _ := io.ReadAll(resp.Body)
  731. return fmt.Errorf("registry responded with code %d: %v", resp.StatusCode, string(body))
  732. }
  733. return nil
  734. }
  735. func downloadBlob(mp ModelPath, digest string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
  736. fp, err := GetBlobsPath(digest)
  737. if err != nil {
  738. return err
  739. }
  740. if fi, _ := os.Stat(fp); fi != nil {
  741. // we already have the file, so return
  742. fn(api.ProgressResponse{
  743. Digest: digest,
  744. Total: int(fi.Size()),
  745. Completed: int(fi.Size()),
  746. })
  747. return nil
  748. }
  749. var size int64
  750. fi, err := os.Stat(fp + "-partial")
  751. switch {
  752. case errors.Is(err, os.ErrNotExist):
  753. // noop, file doesn't exist so create it
  754. case err != nil:
  755. return fmt.Errorf("stat: %w", err)
  756. default:
  757. size = fi.Size()
  758. }
  759. url := fmt.Sprintf("%s/v2/%s/blobs/%s", mp.Registry, mp.GetNamespaceRepository(), digest)
  760. headers := map[string]string{
  761. "Range": fmt.Sprintf("bytes=%d-", size),
  762. }
  763. resp, err := makeRequest("GET", url, headers, nil, regOpts)
  764. if err != nil {
  765. log.Printf("couldn't download blob: %v", err)
  766. return err
  767. }
  768. defer resp.Body.Close()
  769. if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusPartialContent {
  770. body, _ := ioutil.ReadAll(resp.Body)
  771. return fmt.Errorf("registry responded with code %d: %v", resp.StatusCode, string(body))
  772. }
  773. err = os.MkdirAll(path.Dir(fp), 0o700)
  774. if err != nil {
  775. return fmt.Errorf("make blobs directory: %w", err)
  776. }
  777. out, err := os.OpenFile(fp+"-partial", os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644)
  778. if err != nil {
  779. panic(err)
  780. }
  781. defer out.Close()
  782. remaining, _ := strconv.ParseInt(resp.Header.Get("Content-Length"), 10, 64)
  783. completed := size
  784. total := remaining + completed
  785. for {
  786. fn(api.ProgressResponse{
  787. Status: fmt.Sprintf("downloading %s", digest),
  788. Digest: digest,
  789. Total: int(total),
  790. Completed: int(completed),
  791. })
  792. if completed >= total {
  793. if err := out.Close(); err != nil {
  794. return err
  795. }
  796. if err := os.Rename(fp+"-partial", fp); err != nil {
  797. fn(api.ProgressResponse{
  798. Status: fmt.Sprintf("error renaming file: %v", err),
  799. Digest: digest,
  800. Total: int(total),
  801. Completed: int(completed),
  802. })
  803. return err
  804. }
  805. break
  806. }
  807. n, err := io.CopyN(out, resp.Body, 8192)
  808. if err != nil && !errors.Is(err, io.EOF) {
  809. return err
  810. }
  811. completed += n
  812. }
  813. log.Printf("success getting %s\n", digest)
  814. return nil
  815. }
  816. func makeRequest(method, url string, headers map[string]string, body io.Reader, regOpts *RegistryOptions) (*http.Response, error) {
  817. if !strings.HasPrefix(url, "http") {
  818. if regOpts.Insecure {
  819. url = "http://" + url
  820. } else {
  821. url = "https://" + url
  822. }
  823. }
  824. log.Printf("url = %s", url)
  825. req, err := http.NewRequest(method, url, body)
  826. if err != nil {
  827. return nil, err
  828. }
  829. for k, v := range headers {
  830. req.Header.Set(k, v)
  831. }
  832. // TODO: better auth
  833. if regOpts.Username != "" && regOpts.Password != "" {
  834. req.SetBasicAuth(regOpts.Username, regOpts.Password)
  835. }
  836. client := &http.Client{
  837. CheckRedirect: func(req *http.Request, via []*http.Request) error {
  838. if len(via) >= 10 {
  839. return fmt.Errorf("too many redirects")
  840. }
  841. log.Printf("redirected to: %s\n", req.URL)
  842. return nil
  843. },
  844. }
  845. resp, err := client.Do(req)
  846. if err != nil {
  847. return nil, err
  848. }
  849. return resp, nil
  850. }
  851. func verifyBlob(digest string) error {
  852. fp, err := GetBlobsPath(digest)
  853. if err != nil {
  854. return err
  855. }
  856. f, err := os.Open(fp)
  857. if err != nil {
  858. return err
  859. }
  860. defer f.Close()
  861. fileDigest, _ := GetSHA256Digest(f)
  862. if digest != fileDigest {
  863. return fmt.Errorf("digest mismatch: want %s, got %s", digest, fileDigest)
  864. }
  865. return nil
  866. }