images.go 24 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054
  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(request.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 {
  94. return nil, err
  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) error {
  420. mp := ParseModelPath(name)
  421. manifest, err := GetManifest(mp)
  422. if err != nil {
  423. return err
  424. }
  425. deleteMap := make(map[string]bool)
  426. for _, layer := range manifest.Layers {
  427. deleteMap[layer.Digest] = true
  428. }
  429. deleteMap[manifest.Config.Digest] = true
  430. fp, err := GetManifestPath()
  431. if err != nil {
  432. return err
  433. }
  434. err = filepath.Walk(fp, func(path string, info os.FileInfo, err error) error {
  435. if err != nil {
  436. return err
  437. }
  438. if !info.IsDir() {
  439. path := path[len(fp)+1:]
  440. slashIndex := strings.LastIndex(path, "/")
  441. if slashIndex == -1 {
  442. return nil
  443. }
  444. tag := path[:slashIndex] + ":" + path[slashIndex+1:]
  445. fmp := ParseModelPath(tag)
  446. // skip the manifest we're trying to delete
  447. if mp.GetFullTagname() == fmp.GetFullTagname() {
  448. return nil
  449. }
  450. // save (i.e. delete from the deleteMap) any files used in other manifests
  451. manifest, err := GetManifest(fmp)
  452. if err != nil {
  453. log.Printf("skipping file: %s", fp)
  454. return nil
  455. }
  456. for _, layer := range manifest.Layers {
  457. delete(deleteMap, layer.Digest)
  458. }
  459. delete(deleteMap, manifest.Config.Digest)
  460. }
  461. return nil
  462. })
  463. // only delete the files which are still in the deleteMap
  464. for k, v := range deleteMap {
  465. if v {
  466. fp, err := GetBlobsPath(k)
  467. if err != nil {
  468. log.Printf("couldn't get file path for '%s': %v", k, err)
  469. continue
  470. }
  471. if err := os.Remove(fp); err != nil {
  472. log.Printf("couldn't remove file '%s': %v", fp, err)
  473. continue
  474. }
  475. }
  476. }
  477. fp, err = mp.GetManifestPath(false)
  478. if err != nil {
  479. return err
  480. }
  481. err = os.Remove(fp)
  482. if err != nil {
  483. log.Printf("couldn't remove manifest file '%s': %v", fp, err)
  484. return err
  485. }
  486. return nil
  487. }
  488. func PushModel(name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
  489. mp := ParseModelPath(name)
  490. fn(api.ProgressResponse{Status: "retrieving manifest"})
  491. manifest, err := GetManifest(mp)
  492. if err != nil {
  493. fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
  494. return err
  495. }
  496. var layers []*Layer
  497. for _, layer := range manifest.Layers {
  498. layers = append(layers, layer)
  499. }
  500. layers = append(layers, &manifest.Config)
  501. for _, layer := range layers {
  502. exists, err := checkBlobExistence(mp, layer.Digest, regOpts)
  503. if err != nil {
  504. return err
  505. }
  506. if exists {
  507. fn(api.ProgressResponse{
  508. Status: "using existing layer",
  509. Digest: layer.Digest,
  510. Total: layer.Size,
  511. Completed: layer.Size,
  512. })
  513. log.Printf("Layer %s already exists", layer.Digest)
  514. continue
  515. }
  516. fn(api.ProgressResponse{
  517. Status: "starting upload",
  518. Digest: layer.Digest,
  519. Total: layer.Size,
  520. })
  521. location, err := startUpload(mp, regOpts)
  522. if err != nil {
  523. log.Printf("couldn't start upload: %v", err)
  524. return err
  525. }
  526. err = uploadBlobChunked(mp, location, layer, regOpts, fn)
  527. if err != nil {
  528. log.Printf("error uploading blob: %v", err)
  529. return err
  530. }
  531. }
  532. fn(api.ProgressResponse{Status: "pushing manifest"})
  533. url := fmt.Sprintf("%s/v2/%s/manifests/%s", mp.Registry, mp.GetNamespaceRepository(), mp.Tag)
  534. headers := map[string]string{
  535. "Content-Type": "application/vnd.docker.distribution.manifest.v2+json",
  536. }
  537. manifestJSON, err := json.Marshal(manifest)
  538. if err != nil {
  539. return err
  540. }
  541. resp, err := makeRequest("PUT", url, headers, bytes.NewReader(manifestJSON), regOpts)
  542. if err != nil {
  543. return err
  544. }
  545. defer resp.Body.Close()
  546. // Check for success: For a successful upload, the Docker registry will respond with a 201 Created
  547. if resp.StatusCode != http.StatusCreated {
  548. body, _ := io.ReadAll(resp.Body)
  549. return fmt.Errorf("registry responded with code %d: %v", resp.StatusCode, string(body))
  550. }
  551. fn(api.ProgressResponse{Status: "success"})
  552. return nil
  553. }
  554. func PullModel(name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
  555. mp := ParseModelPath(name)
  556. fn(api.ProgressResponse{Status: "pulling manifest"})
  557. manifest, err := pullModelManifest(mp, regOpts)
  558. if err != nil {
  559. return fmt.Errorf("pull model manifest: %q", err)
  560. }
  561. var layers []*Layer
  562. layers = append(layers, manifest.Layers...)
  563. layers = append(layers, &manifest.Config)
  564. for _, layer := range layers {
  565. if err := downloadBlob(mp, layer.Digest, regOpts, fn); err != nil {
  566. return err
  567. }
  568. }
  569. fn(api.ProgressResponse{Status: "verifying sha256 digest"})
  570. for _, layer := range layers {
  571. if err := verifyBlob(layer.Digest); err != nil {
  572. return err
  573. }
  574. }
  575. fn(api.ProgressResponse{Status: "writing manifest"})
  576. manifestJSON, err := json.Marshal(manifest)
  577. if err != nil {
  578. return err
  579. }
  580. fp, err := mp.GetManifestPath(true)
  581. if err != nil {
  582. return err
  583. }
  584. err = os.WriteFile(fp, manifestJSON, 0o644)
  585. if err != nil {
  586. log.Printf("couldn't write to %s", fp)
  587. return err
  588. }
  589. fn(api.ProgressResponse{Status: "success"})
  590. return nil
  591. }
  592. func pullModelManifest(mp ModelPath, regOpts *RegistryOptions) (*ManifestV2, error) {
  593. url := fmt.Sprintf("%s/v2/%s/manifests/%s", mp.Registry, mp.GetNamespaceRepository(), mp.Tag)
  594. headers := map[string]string{
  595. "Accept": "application/vnd.docker.distribution.manifest.v2+json",
  596. }
  597. resp, err := makeRequest("GET", url, headers, nil, regOpts)
  598. if err != nil {
  599. log.Printf("couldn't get manifest: %v", err)
  600. return nil, err
  601. }
  602. defer resp.Body.Close()
  603. // Check for success: For a successful upload, the Docker registry will respond with a 201 Created
  604. if resp.StatusCode != http.StatusOK {
  605. body, _ := io.ReadAll(resp.Body)
  606. return nil, fmt.Errorf("registry responded with code %d: %s", resp.StatusCode, body)
  607. }
  608. var m *ManifestV2
  609. if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
  610. return nil, err
  611. }
  612. return m, err
  613. }
  614. func createConfigLayer(layers []string) (*LayerReader, error) {
  615. // TODO change architecture and OS
  616. config := ConfigV2{
  617. Architecture: "arm64",
  618. OS: "linux",
  619. RootFS: RootFS{
  620. Type: "layers",
  621. DiffIDs: layers,
  622. },
  623. }
  624. configJSON, err := json.Marshal(config)
  625. if err != nil {
  626. return nil, err
  627. }
  628. digest, size := GetSHA256Digest(bytes.NewBuffer(configJSON))
  629. layer := &LayerReader{
  630. Layer: Layer{
  631. MediaType: "application/vnd.docker.container.image.v1+json",
  632. Digest: digest,
  633. Size: size,
  634. },
  635. Reader: bytes.NewBuffer(configJSON),
  636. }
  637. return layer, nil
  638. }
  639. // GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer
  640. func GetSHA256Digest(r io.Reader) (string, int) {
  641. h := sha256.New()
  642. n, err := io.Copy(h, r)
  643. if err != nil {
  644. log.Fatal(err)
  645. }
  646. return fmt.Sprintf("sha256:%x", h.Sum(nil)), int(n)
  647. }
  648. func startUpload(mp ModelPath, regOpts *RegistryOptions) (string, error) {
  649. url := fmt.Sprintf("%s/v2/%s/blobs/uploads/", mp.Registry, mp.GetNamespaceRepository())
  650. resp, err := makeRequest("POST", url, nil, nil, regOpts)
  651. if err != nil {
  652. log.Printf("couldn't start upload: %v", err)
  653. return "", err
  654. }
  655. defer resp.Body.Close()
  656. // Check for success
  657. if resp.StatusCode != http.StatusAccepted {
  658. body, _ := io.ReadAll(resp.Body)
  659. return "", fmt.Errorf("registry responded with code %d: %s", resp.StatusCode, body)
  660. }
  661. // Extract UUID location from header
  662. location := resp.Header.Get("Location")
  663. if location == "" {
  664. return "", fmt.Errorf("location header is missing in response")
  665. }
  666. return location, nil
  667. }
  668. // Function to check if a blob already exists in the Docker registry
  669. func checkBlobExistence(mp ModelPath, digest string, regOpts *RegistryOptions) (bool, error) {
  670. url := fmt.Sprintf("%s/v2/%s/blobs/%s", mp.Registry, mp.GetNamespaceRepository(), digest)
  671. resp, err := makeRequest("HEAD", url, nil, nil, regOpts)
  672. if err != nil {
  673. log.Printf("couldn't check for blob: %v", err)
  674. return false, err
  675. }
  676. defer resp.Body.Close()
  677. // Check for success: If the blob exists, the Docker registry will respond with a 200 OK
  678. return resp.StatusCode == http.StatusOK, nil
  679. }
  680. func uploadBlobChunked(mp ModelPath, location string, layer *Layer, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
  681. // TODO allow resumability
  682. // TODO allow canceling uploads via DELETE
  683. // TODO allow cross repo blob mount
  684. // Create URL
  685. url := fmt.Sprintf("%s", location)
  686. fp, err := GetBlobsPath(layer.Digest)
  687. if err != nil {
  688. return err
  689. }
  690. f, err := os.Open(fp)
  691. if err != nil {
  692. return err
  693. }
  694. headers := make(map[string]string)
  695. headers["Content-Type"] = "application/octet-stream"
  696. chunkSize := 1 << 20
  697. buf := make([]byte, chunkSize)
  698. var totalUploaded int
  699. for {
  700. n, err := f.Read(buf)
  701. if err != nil {
  702. return err
  703. }
  704. headers["Content-Length"] = fmt.Sprintf("%d", n)
  705. headers["Content-Range"] = fmt.Sprintf("%d-%d", totalUploaded, totalUploaded+n-1)
  706. fn(api.ProgressResponse{
  707. Status: fmt.Sprintf("uploading %s", layer.Digest),
  708. Digest: layer.Digest,
  709. Total: int(layer.Size),
  710. Completed: int(totalUploaded),
  711. })
  712. // change the buffersize for the last chunk
  713. if n < chunkSize {
  714. buf = buf[:n]
  715. }
  716. resp, err := makeRequest("PATCH", url, headers, bytes.NewReader(buf), regOpts)
  717. if err != nil {
  718. log.Printf("couldn't upload blob: %v", err)
  719. return err
  720. }
  721. defer resp.Body.Close()
  722. url = resp.Header.Get("Location")
  723. // Check for success: For a successful upload, the Docker registry will respond with a 201 Created
  724. if resp.StatusCode != http.StatusAccepted {
  725. fn(api.ProgressResponse{
  726. Status: fmt.Sprintf("error uploading layer"),
  727. Digest: layer.Digest,
  728. Total: int(layer.Size),
  729. Completed: int(totalUploaded),
  730. })
  731. body, _ := io.ReadAll(resp.Body)
  732. return fmt.Errorf("registry responded with code %d: %v", resp.StatusCode, string(body))
  733. }
  734. totalUploaded += n
  735. if totalUploaded >= layer.Size {
  736. url = fmt.Sprintf("%s&digest=%s", url, layer.Digest)
  737. // finish the upload
  738. resp, err := makeRequest("PUT", url, nil, nil, regOpts)
  739. if err != nil {
  740. log.Printf("couldn't finish upload: %v", err)
  741. return err
  742. }
  743. defer resp.Body.Close()
  744. if resp.StatusCode != http.StatusCreated {
  745. body, _ := io.ReadAll(resp.Body)
  746. return fmt.Errorf("registry responded with code %d: %v", resp.StatusCode, string(body))
  747. }
  748. break
  749. }
  750. }
  751. return nil
  752. }
  753. func downloadBlob(mp ModelPath, digest string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
  754. fp, err := GetBlobsPath(digest)
  755. if err != nil {
  756. return err
  757. }
  758. if fi, _ := os.Stat(fp); fi != nil {
  759. // we already have the file, so return
  760. fn(api.ProgressResponse{
  761. Digest: digest,
  762. Total: int(fi.Size()),
  763. Completed: int(fi.Size()),
  764. })
  765. return nil
  766. }
  767. var size int64
  768. fi, err := os.Stat(fp + "-partial")
  769. switch {
  770. case errors.Is(err, os.ErrNotExist):
  771. // noop, file doesn't exist so create it
  772. case err != nil:
  773. return fmt.Errorf("stat: %w", err)
  774. default:
  775. size = fi.Size()
  776. }
  777. url := fmt.Sprintf("%s/v2/%s/blobs/%s", mp.Registry, mp.GetNamespaceRepository(), digest)
  778. headers := map[string]string{
  779. "Range": fmt.Sprintf("bytes=%d-", size),
  780. }
  781. resp, err := makeRequest("GET", url, headers, nil, regOpts)
  782. if err != nil {
  783. log.Printf("couldn't download blob: %v", err)
  784. return err
  785. }
  786. defer resp.Body.Close()
  787. if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusPartialContent {
  788. body, _ := ioutil.ReadAll(resp.Body)
  789. return fmt.Errorf("registry responded with code %d: %v", resp.StatusCode, string(body))
  790. }
  791. err = os.MkdirAll(path.Dir(fp), 0o700)
  792. if err != nil {
  793. return fmt.Errorf("make blobs directory: %w", err)
  794. }
  795. out, err := os.OpenFile(fp+"-partial", os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644)
  796. if err != nil {
  797. panic(err)
  798. }
  799. defer out.Close()
  800. remaining, _ := strconv.ParseInt(resp.Header.Get("Content-Length"), 10, 64)
  801. completed := size
  802. total := remaining + completed
  803. for {
  804. fn(api.ProgressResponse{
  805. Status: fmt.Sprintf("downloading %s", digest),
  806. Digest: digest,
  807. Total: int(total),
  808. Completed: int(completed),
  809. })
  810. if completed >= total {
  811. if err := out.Close(); err != nil {
  812. return err
  813. }
  814. if err := os.Rename(fp+"-partial", fp); err != nil {
  815. fn(api.ProgressResponse{
  816. Status: fmt.Sprintf("error renaming file: %v", err),
  817. Digest: digest,
  818. Total: int(total),
  819. Completed: int(completed),
  820. })
  821. return err
  822. }
  823. break
  824. }
  825. n, err := io.CopyN(out, resp.Body, 8192)
  826. if err != nil && !errors.Is(err, io.EOF) {
  827. return err
  828. }
  829. completed += n
  830. }
  831. log.Printf("success getting %s\n", digest)
  832. return nil
  833. }
  834. func makeRequest(method, url string, headers map[string]string, body io.Reader, regOpts *RegistryOptions) (*http.Response, error) {
  835. if !strings.HasPrefix(url, "http") {
  836. if regOpts.Insecure {
  837. url = "http://" + url
  838. } else {
  839. url = "https://" + url
  840. }
  841. }
  842. req, err := http.NewRequest(method, url, body)
  843. if err != nil {
  844. return nil, err
  845. }
  846. for k, v := range headers {
  847. req.Header.Set(k, v)
  848. }
  849. // TODO: better auth
  850. if regOpts.Username != "" && regOpts.Password != "" {
  851. req.SetBasicAuth(regOpts.Username, regOpts.Password)
  852. }
  853. client := &http.Client{
  854. CheckRedirect: func(req *http.Request, via []*http.Request) error {
  855. if len(via) >= 10 {
  856. return fmt.Errorf("too many redirects")
  857. }
  858. log.Printf("redirected to: %s\n", req.URL)
  859. return nil
  860. },
  861. }
  862. resp, err := client.Do(req)
  863. if err != nil {
  864. return nil, err
  865. }
  866. return resp, nil
  867. }
  868. func verifyBlob(digest string) error {
  869. fp, err := GetBlobsPath(digest)
  870. if err != nil {
  871. return err
  872. }
  873. f, err := os.Open(fp)
  874. if err != nil {
  875. return err
  876. }
  877. defer f.Close()
  878. fileDigest, _ := GetSHA256Digest(f)
  879. if digest != fileDigest {
  880. return fmt.Errorf("digest mismatch: want %s, got %s", digest, fileDigest)
  881. }
  882. return nil
  883. }