diff --git a/__tests__/context.test.ts b/__tests__/context.test.ts
index 5e813ac..0f0f79e 100644
--- a/__tests__/context.test.ts
+++ b/__tests__/context.test.ts
@@ -21,8 +21,11 @@ describe('getInputs', () => {
         ['set-host', 'false'],
       ]),
       {
-        version: 'v24.0.8',
-        channel: '',
+        source: {
+          type: 'archive',
+          version: 'v24.0.8',
+          channel: 'stable'
+        },
         context: '',
         daemonConfig: '',
         setHost: false
@@ -38,8 +41,11 @@ describe('getInputs', () => {
         ['set-host', 'false'],
       ]),
       {
-        version: 'v24.0.0-rc.4',
-        channel: 'test',
+        source: {
+          type: 'archive',
+          version: 'v24.0.0-rc.4',
+          channel: 'test'
+        },
         context: 'foo',
         daemonConfig: `{"debug":true,"features":{"containerd-snapshotter":true}}`,
         setHost: false
@@ -51,13 +57,100 @@ describe('getInputs', () => {
         ['set-host', 'true'],
       ]),
       {
-        version: 'latest',
-        channel: '',
+        source: {
+          type: 'archive',
+          version: 'latest',
+          channel: 'stable',
+        },
         context: '',
         daemonConfig: '',
         setHost: true
       } as context.Inputs
-    ]
+    ],
+    [
+    3,
+      new Map<string, string>([
+        ['version', 'type=image,tag=master'],
+        ['context', 'foo'],
+        ['daemon-config', `{"debug":true,"features":{"containerd-snapshotter":true}}`],
+        ['set-host', 'false'],
+      ]),
+      {
+        source: {
+          type: 'image',
+          tag: 'master',
+        },
+        context: 'foo',
+        daemonConfig: `{"debug":true,"features":{"containerd-snapshotter":true}}`,
+        setHost: false
+      } as context.Inputs
+    ],
+    [
+      4,
+      new Map<string, string>([
+        ['version', 'type=image'],
+        ['set-host', 'false'],
+      ]),
+      {
+        source: {
+          type: 'image',
+          tag: 'latest',
+        },
+        context: '',
+        daemonConfig: '',
+        setHost: false
+      } as context.Inputs
+    ],
+    [
+      5,
+      new Map<string, string>([
+        ['version', 'type=archive'],
+        ['set-host', 'false'],
+      ]),
+      {
+        source: {
+          type: 'archive',
+          version: 'latest',
+          channel: 'stable',
+        },
+        setHost: false,
+        context: '',
+        daemonConfig: '',
+      } as context.Inputs
+    ],
+    [
+      6,
+      new Map<string, string>([
+        ['version', 'version=v27.2.0,channel=test'],
+        ['set-host', 'false'],
+      ]),
+      {
+        source: {
+          type: 'archive',
+          version: 'v27.2.0',
+          channel: 'test',
+        },
+        setHost: false,
+        context: '',
+        daemonConfig: '',
+      } as context.Inputs
+    ],
+    [
+      7,
+      new Map<string, string>([
+        ['version', 'type=image,tag=27.2.1'],
+        ['set-host', 'false'],
+      ]),
+      {
+        source: {
+          type: 'image',
+          tag: '27.2.1',
+        },
+        setHost: false,
+        context: '',
+        daemonConfig: '',
+      } as context.Inputs
+    ],
   ])(
     '[%d] given %p as inputs, returns %p',
     async (num: number, inputs: Map<string, string>, expected: context.Inputs) => {
diff --git a/src/context.ts b/src/context.ts
index 34ba9a7..b416a63 100644
--- a/src/context.ts
+++ b/src/context.ts
@@ -1,19 +1,92 @@
 import * as core from '@actions/core';
+import {InstallSource} from '@docker/actions-toolkit/lib/docker/install';
+import { parse } from 'csv-parse';
 
 export interface Inputs {
-  version: string;
-  channel: string;
+  source: InstallSource;
   daemonConfig?: string;
   context: string;
   setHost: boolean;
 }
 
 export function getInputs(): Inputs {
+  const rawVersion = core.getInput('version') || 'latest';
+  const source = parseSource(rawVersion);
+  const channel = core.getInput('channel');
+  if (channel && source.type === 'archive') {
+    source.channel = channel;
+  }
+
   return {
-    version: core.getInput('version') || 'latest',
-    channel: core.getInput('channel'),
+    source: source,
     daemonConfig: core.getInput('daemon-config'),
     context: core.getInput('context'),
     setHost: core.getBooleanInput('set-host')
   };
 }
+
+function parseSource(input: string): InstallSource {
+  let [type, version, channel, tag] = ['archive', input, 'stable', 'latest'];
+
+  const fields = parse(version, {
+    relaxColumnCount: true,
+    skipEmptyLines: true
+  })[0];
+  for (const field of fields) {
+    const parts = field
+      .toString()
+      .split(/(?<=^[^=]+?)=/)
+      .map(item => item.trim());
+
+    switch (parts[0]) {
+      case 'type':
+        type = parts[1];
+        break;
+      case 'version':
+        version = parts[1];
+        break;
+      case 'channel':
+        channel = parts[1];
+        break;
+      case 'tag':
+        tag = parts[1];
+        break;
+      default:
+        throw new Error(`Invalid field: ${parts[0]}`);
+    }
+  }
+
+  if (!type) {
+    throw new Error(`Invalid type: ${type}`);
+  }
+  if (!channel) {
+    throw new Error(`Invalid channel: ${channel}`);
+  }
+  if (!version) {
+    throw new Error(`Invalid version: ${version}`);
+  }
+  if (!tag) {
+    throw new Error(`Invalid tag: ${tag}`);
+  }
+
+  let src: InstallSource;
+  switch (type) {
+    case 'archive':
+      src = {
+        type: 'archive',
+        version: version,
+        channel: channel
+      };
+      break;
+    case 'image':
+      src = {
+        type: 'image',
+        tag: tag
+      };
+      break;
+    default:
+      throw new Error(`Invalid version: ${input}`);
+  }
+
+  return src;
+}
diff --git a/src/main.ts b/src/main.ts
index 5b53745..1b9303e 100644
--- a/src/main.ts
+++ b/src/main.ts
@@ -21,16 +21,12 @@ actionsToolkit.run(
 
     const install = new Install({
       runDir: runDir,
-      source: {
-        type: 'archive',
-        version: input.version,
-        channel: input.channel || 'stable'
-      },
+      source: input.source,
       contextName: input.context || 'setup-docker-action',
       daemonConfig: input.daemonConfig
     });
     let toolDir;
-    if (!(await Docker.isAvailable()) || input.version) {
+    if (!(await Docker.isAvailable()) || input.source) {
       await core.group(`Download docker`, async () => {
         toolDir = await install.download();
       });